831 lines
29 KiB
Text
831 lines
29 KiB
Text
/-
|
||
Copyright (c) 2023 Lean FRO, LLC. All rights reserved.
|
||
Released under Apache 2.0 license as described in the file LICENSE.
|
||
Authors: Joe Hendrix, Scott Morrison
|
||
-/
|
||
prelude
|
||
import Lean.Meta.CompletionName
|
||
import Lean.Meta.DiscrTree
|
||
|
||
/-!
|
||
# Lazy Discrimination Tree
|
||
|
||
This file defines a new type of discrimination tree optimized for rapid
|
||
population of imported modules for use in tactics. It uses a lazy
|
||
initialization strategy.
|
||
|
||
The discrimination tree can be created through
|
||
`createImportedEnvironment`. This creates a discrimination tree from all
|
||
imported modules in an environment using a callback that provides the
|
||
entries as `InitEntry` values.
|
||
|
||
The function `getMatch` can be used to get the values that match the
|
||
expression as well as an updated lazy discrimination tree that has
|
||
elaborated additional parts of the tree.
|
||
-/
|
||
namespace Lean.Meta.LazyDiscrTree
|
||
|
||
-- This namespace contains definitions copied from Lean.Meta.DiscrTree.
|
||
namespace MatchClone
|
||
|
||
/--
|
||
Discrimination tree key.
|
||
-/
|
||
private inductive Key where
|
||
| const : Name → Nat → Key
|
||
| fvar : FVarId → Nat → Key
|
||
| lit : Literal → Key
|
||
| star : Key
|
||
| other : Key
|
||
| arrow : Key
|
||
| proj : Name → Nat → Nat → Key
|
||
deriving Inhabited, BEq, Repr
|
||
|
||
namespace Key
|
||
|
||
/-- Hash function -/
|
||
protected def hash : Key → UInt64
|
||
| .const n a => mixHash 5237 $ mixHash n.hash (hash a)
|
||
| .fvar n a => mixHash 3541 $ mixHash (hash n) (hash a)
|
||
| .lit v => mixHash 1879 $ hash v
|
||
| .star => 7883
|
||
| .other => 2411
|
||
| .arrow => 17
|
||
| .proj s i a => mixHash (hash a) $ mixHash (hash s) (hash i)
|
||
|
||
instance : Hashable Key := ⟨Key.hash⟩
|
||
|
||
end Key
|
||
|
||
private def tmpMVarId : MVarId := { name := `_discr_tree_tmp }
|
||
private def tmpStar := mkMVar tmpMVarId
|
||
|
||
/--
|
||
Returns true iff the argument should be treated as a "wildcard" by the
|
||
discrimination tree.
|
||
|
||
This includes proofs, instance implicit arguments, implicit arguments,
|
||
and terms of the form `noIndexing t`
|
||
|
||
This is a clone of `Lean.Meta.DiscrTree.ignoreArg` and mainly added to
|
||
avoid coupling between `DiscrTree` and `LazyDiscrTree` while both are
|
||
potentially subject to independent changes.
|
||
-/
|
||
private def ignoreArg (a : Expr) (i : Nat) (infos : Array ParamInfo) : MetaM Bool := do
|
||
if h : i < infos.size then
|
||
let info := infos.get ⟨i, h⟩
|
||
if info.isInstImplicit then
|
||
return true
|
||
else if info.isImplicit || info.isStrictImplicit then
|
||
return not (← isType a)
|
||
else
|
||
isProof a
|
||
else
|
||
isProof a
|
||
|
||
private partial def pushArgsAux (infos : Array ParamInfo) : Nat → Expr → Array Expr → MetaM (Array Expr)
|
||
| i, .app f a, todo => do
|
||
if (← ignoreArg a i infos) then
|
||
pushArgsAux infos (i-1) f (todo.push tmpStar)
|
||
else
|
||
pushArgsAux infos (i-1) f (todo.push a)
|
||
| _, _, todo => return todo
|
||
|
||
/--
|
||
Returns `true` if `e` is one of the following
|
||
- A nat literal (numeral)
|
||
- `Nat.zero`
|
||
- `Nat.succ x` where `isNumeral x`
|
||
- `OfNat.ofNat _ x _` where `isNumeral x` -/
|
||
private partial def isNumeral (e : Expr) : Bool :=
|
||
if e.isRawNatLit then true
|
||
else
|
||
let f := e.getAppFn
|
||
if !f.isConst then false
|
||
else
|
||
let fName := f.constName!
|
||
if fName == ``Nat.succ && e.getAppNumArgs == 1 then isNumeral e.appArg!
|
||
else if fName == ``OfNat.ofNat && e.getAppNumArgs == 3 then isNumeral (e.getArg! 1)
|
||
else if fName == ``Nat.zero && e.getAppNumArgs == 0 then true
|
||
else false
|
||
|
||
private partial def toNatLit? (e : Expr) : Option Literal :=
|
||
if isNumeral e then
|
||
if let some n := loop e then
|
||
some (.natVal n)
|
||
else
|
||
none
|
||
else
|
||
none
|
||
where
|
||
loop (e : Expr) : OptionT Id Nat := do
|
||
let f := e.getAppFn
|
||
match f with
|
||
| .lit (.natVal n) => return n
|
||
| .const fName .. =>
|
||
if fName == ``Nat.succ && e.getAppNumArgs == 1 then
|
||
let r ← loop e.appArg!
|
||
return r+1
|
||
else if fName == ``OfNat.ofNat && e.getAppNumArgs == 3 then
|
||
loop (e.getArg! 1)
|
||
else if fName == ``Nat.zero && e.getAppNumArgs == 0 then
|
||
return 0
|
||
else
|
||
failure
|
||
| _ => failure
|
||
|
||
private def isNatType (e : Expr) : MetaM Bool :=
|
||
return (← whnf e).isConstOf ``Nat
|
||
|
||
/--
|
||
Returns `true` if `e` is one of the following
|
||
- `Nat.add _ k` where `isNumeral k`
|
||
- `Add.add Nat _ _ k` where `isNumeral k`
|
||
- `HAdd.hAdd _ Nat _ _ k` where `isNumeral k`
|
||
- `Nat.succ _`
|
||
This function assumes `e.isAppOf fName`
|
||
-/
|
||
private def isNatOffset (fName : Name) (e : Expr) : MetaM Bool := do
|
||
if fName == ``Nat.add && e.getAppNumArgs == 2 then
|
||
return isNumeral e.appArg!
|
||
else if fName == ``Add.add && e.getAppNumArgs == 4 then
|
||
if (← isNatType (e.getArg! 0)) then return isNumeral e.appArg! else return false
|
||
else if fName == ``HAdd.hAdd && e.getAppNumArgs == 6 then
|
||
if (← isNatType (e.getArg! 1)) then return isNumeral e.appArg! else return false
|
||
else
|
||
return fName == ``Nat.succ && e.getAppNumArgs == 1
|
||
|
||
/-
|
||
This is a hook to determine if we should add an expression as a wildcard pattern.
|
||
|
||
Clone of `Lean.Meta.DiscrTree.shouldAddAsStar`. See it for more discussion.
|
||
-/
|
||
private def shouldAddAsStar (fName : Name) (e : Expr) : MetaM Bool := do
|
||
isNatOffset fName e
|
||
|
||
/--
|
||
Eliminate loose bound variables via beta-reduction.
|
||
|
||
This is primarily used to reduce pi-terms `∀(x : P), T` into
|
||
non-dependend functions `P → T`. The latter has a more specific
|
||
discrimination tree key `.arrow..` and this improves the accuracy of the
|
||
discrimination tree.
|
||
|
||
Clone of `Lean.Meta.DiscrTree.elimLooseBVarsByBeta`. See it for more
|
||
discussion.
|
||
-/
|
||
private def elimLooseBVarsByBeta (e : Expr) : CoreM Expr :=
|
||
Core.transform e
|
||
(pre := fun e => do
|
||
if !e.hasLooseBVars then
|
||
return .done e
|
||
else if e.isHeadBetaTarget then
|
||
return .visit e.headBeta
|
||
else
|
||
return .continue)
|
||
|
||
private def getKeyArgs (e : Expr) (isMatch root : Bool) (config : WhnfCoreConfig) :
|
||
MetaM (Key × Array Expr) := do
|
||
let e ← DiscrTree.reduceDT e root config
|
||
unless root do
|
||
-- See pushArgs
|
||
if let some v := toNatLit? e then
|
||
return (.lit v, #[])
|
||
match e.getAppFn with
|
||
| .lit v => return (.lit v, #[])
|
||
| .const c _ =>
|
||
if (← getConfig).isDefEqStuckEx && e.hasExprMVar then
|
||
if (← isReducible c) then
|
||
/- `e` is a term `c ...` s.t. `c` is reducible and `e` has metavariables, but it was not
|
||
unfolded. This can happen if the metavariables in `e` are "blocking" smart unfolding.
|
||
If `isDefEqStuckEx` is enabled, then we must throw the `isDefEqStuck` exception to
|
||
postpone TC resolution.
|
||
-/
|
||
Meta.throwIsDefEqStuck
|
||
else if let some matcherInfo := isMatcherAppCore? (← getEnv) e then
|
||
-- A matcher application is stuck if one of the discriminants has a metavariable
|
||
let args := e.getAppArgs
|
||
let start := matcherInfo.getFirstDiscrPos
|
||
for arg in args[ start : start + matcherInfo.numDiscrs ] do
|
||
if arg.hasExprMVar then
|
||
Meta.throwIsDefEqStuck
|
||
else if (← isRec c) then
|
||
/- Similar to the previous case, but for `match` and recursor applications. It may be stuck
|
||
(i.e., did not reduce) because of metavariables. -/
|
||
Meta.throwIsDefEqStuck
|
||
let nargs := e.getAppNumArgs
|
||
return (.const c nargs, e.getAppRevArgs)
|
||
| .fvar fvarId =>
|
||
let nargs := e.getAppNumArgs
|
||
return (.fvar fvarId nargs, e.getAppRevArgs)
|
||
| .mvar mvarId =>
|
||
if isMatch then
|
||
return (.other, #[])
|
||
else do
|
||
let ctx ← read
|
||
if ctx.config.isDefEqStuckEx then
|
||
/-
|
||
When the configuration flag `isDefEqStuckEx` is set to true,
|
||
we want `isDefEq` to throw an exception whenever it tries to assign
|
||
a read-only metavariable.
|
||
This feature is useful for type class resolution where
|
||
we may want to notify the caller that the TC problem may be solvable
|
||
later after it assigns `?m`.
|
||
The method `DiscrTree.getUnify e` returns candidates `c` that may "unify" with `e`.
|
||
That is, `isDefEq c e` may return true. Now, consider `DiscrTree.getUnify d (Add ?m)`
|
||
where `?m` is a read-only metavariable, and the discrimination tree contains the keys
|
||
`HadAdd Nat` and `Add Int`. If `isDefEqStuckEx` is set to true, we must treat `?m` as
|
||
a regular metavariable here, otherwise we return the empty set of candidates.
|
||
This is incorrect because it is equivalent to saying that there is no solution even if
|
||
the caller assigns `?m` and try again. -/
|
||
return (.star, #[])
|
||
else if (← mvarId.isReadOnlyOrSyntheticOpaque) then
|
||
return (.other, #[])
|
||
else
|
||
return (.star, #[])
|
||
| .proj s i a .. =>
|
||
let nargs := e.getAppNumArgs
|
||
return (.proj s i nargs, #[a] ++ e.getAppRevArgs)
|
||
| .forallE _ d b _ =>
|
||
-- See comment at elimLooseBVarsByBeta
|
||
let b ← if b.hasLooseBVars then elimLooseBVarsByBeta b else pure b
|
||
if b.hasLooseBVars then
|
||
return (.other, #[])
|
||
else
|
||
return (.arrow, #[d, b])
|
||
| .bvar _ | .letE _ _ _ _ _ | .lam _ _ _ _ | .mdata _ _ | .app _ _ | .sort _ =>
|
||
return (.other, #[])
|
||
|
||
/-
|
||
Given an expression we are looking for patterns that match, return the key and sub-expressions.
|
||
-/
|
||
private abbrev getMatchKeyArgs (e : Expr) (root : Bool) (config : WhnfCoreConfig) :
|
||
MetaM (Key × Array Expr) :=
|
||
getKeyArgs e (isMatch := true) (root := root) (config := config)
|
||
|
||
end MatchClone
|
||
|
||
export MatchClone (Key Key.const)
|
||
|
||
/--
|
||
An unprocessed entry in the lazy discrimination tree.
|
||
-/
|
||
private abbrev LazyEntry α := Array Expr × ((LocalContext × LocalInstances) × α)
|
||
|
||
/--
|
||
Index identifying trie in a discrimination tree.
|
||
-/
|
||
@[reducible]
|
||
private def TrieIndex := Nat
|
||
|
||
/--
|
||
Discrimination tree trie. See `LazyDiscrTree`.
|
||
-/
|
||
private structure Trie (α : Type) where
|
||
node ::
|
||
/-- Values for matches ending at this trie. -/
|
||
values : Array α
|
||
/-- Index of trie matching star. -/
|
||
star : TrieIndex
|
||
/-- Following matches based on key of trie. -/
|
||
children : HashMap Key TrieIndex
|
||
/-- Lazy entries at this trie that are not processed. -/
|
||
pending : Array (LazyEntry α)
|
||
deriving Inhabited
|
||
|
||
instance : EmptyCollection (Trie α) := ⟨.node #[] 0 {} #[]⟩
|
||
|
||
/-- Push lazy entry to trie. -/
|
||
private def Trie.pushPending : Trie α → LazyEntry α → Trie α
|
||
| .node vs star cs p, e => .node vs star cs (p.push e)
|
||
|
||
end LazyDiscrTree
|
||
|
||
/--
|
||
`LazyDiscrTree` is a variant of the discriminator tree datatype
|
||
`DiscrTree` in Lean core that is designed to be efficiently
|
||
initializable with a large number of patterns. This is useful
|
||
in contexts such as searching an entire Lean environment for
|
||
expressions that match a pattern.
|
||
|
||
Lazy discriminator trees achieve good performance by minimizing
|
||
the amount of work that is done up front to build the discriminator
|
||
tree. When first adding patterns to the tree, only the root
|
||
discriminator key is computed and processing the remaining
|
||
terms is deferred until demanded by a match.
|
||
-/
|
||
structure LazyDiscrTree (α : Type) where
|
||
/-- Configuration for normalization. -/
|
||
config : Lean.Meta.WhnfCoreConfig := {}
|
||
/-- Backing array of trie entries. Should be owned by this trie. -/
|
||
tries : Array (LazyDiscrTree.Trie α) := #[default]
|
||
/-- Map from discriminator trie roots to the index. -/
|
||
roots : Lean.HashMap LazyDiscrTree.Key LazyDiscrTree.TrieIndex := {}
|
||
|
||
namespace LazyDiscrTree
|
||
|
||
open Lean Elab Meta
|
||
|
||
instance : Inhabited (LazyDiscrTree α) where
|
||
default := {}
|
||
|
||
open Lean.Meta.DiscrTree (mkNoindexAnnotation hasNoindexAnnotation reduceDT)
|
||
|
||
/--
|
||
Specialization of Lean.Meta.DiscrTree.pushArgs
|
||
-/
|
||
private def pushArgs (root : Bool) (todo : Array Expr) (e : Expr) (config : WhnfCoreConfig) :
|
||
MetaM (Key × Array Expr) := do
|
||
if hasNoindexAnnotation e then
|
||
return (.star, todo)
|
||
else
|
||
let e ← reduceDT e root config
|
||
let fn := e.getAppFn
|
||
let push (k : Key) (nargs : Nat) (todo : Array Expr) : MetaM (Key × Array Expr) := do
|
||
let info ← getFunInfoNArgs fn nargs
|
||
let todo ← MatchClone.pushArgsAux info.paramInfo (nargs-1) e todo
|
||
return (k, todo)
|
||
match fn with
|
||
| .lit v =>
|
||
return (.lit v, todo)
|
||
| .const c _ =>
|
||
unless root do
|
||
if let some v := MatchClone.toNatLit? e then
|
||
return (.lit v, todo)
|
||
if (← MatchClone.shouldAddAsStar c e) then
|
||
return (.star, todo)
|
||
let nargs := e.getAppNumArgs
|
||
push (.const c nargs) nargs todo
|
||
| .proj s i a =>
|
||
/-
|
||
If `s` is a class, then `a` is an instance. Thus, we annotate `a` with `no_index` since we do
|
||
not index instances. This should only happen if users mark a class projection function as
|
||
`[reducible]`.
|
||
|
||
TODO: add better support for projections that are functions
|
||
-/
|
||
let a := if isClass (← getEnv) s then mkNoindexAnnotation a else a
|
||
let nargs := e.getAppNumArgs
|
||
push (.proj s i nargs) nargs (todo.push a)
|
||
| .fvar _fvarId =>
|
||
return (.star, todo)
|
||
| .mvar mvarId =>
|
||
if mvarId == MatchClone.tmpMVarId then
|
||
-- We use `tmp to mark implicit arguments and proofs
|
||
return (.star, todo)
|
||
else
|
||
failure
|
||
| .forallE _ d b _ =>
|
||
-- See comment at elimLooseBVarsByBeta
|
||
let b ← if b.hasLooseBVars then MatchClone.elimLooseBVarsByBeta b else pure b
|
||
if b.hasLooseBVars then
|
||
return (.other, todo)
|
||
else
|
||
return (.arrow, (todo.push d).push b)
|
||
| _ =>
|
||
return (.other, todo)
|
||
|
||
/-- Initial capacity for key and todo vector. -/
|
||
private def initCapacity := 8
|
||
|
||
/--
|
||
Get the root key and rest of terms of an expression using the specified config.
|
||
-/
|
||
private def rootKey (cfg: WhnfCoreConfig) (e : Expr) : MetaM (Key × Array Expr) :=
|
||
pushArgs true (Array.mkEmpty initCapacity) e cfg
|
||
|
||
private partial def mkPathAux (root : Bool) (todo : Array Expr) (keys : Array Key)
|
||
(config : WhnfCoreConfig) : MetaM (Array Key) := do
|
||
if todo.isEmpty then
|
||
return keys
|
||
else
|
||
let e := todo.back
|
||
let todo := todo.pop
|
||
let (k, todo) ← pushArgs root todo e config
|
||
mkPathAux false todo (keys.push k) config
|
||
|
||
/--
|
||
Create a path from an expression.
|
||
|
||
This differs from Lean.Meta.DiscrTree.mkPath in that the expression
|
||
should uses free variables rather than meta-variables for holes.
|
||
-/
|
||
private def mkPath (e : Expr) (config : WhnfCoreConfig) : MetaM (Array Key) := do
|
||
let todo : Array Expr := .mkEmpty initCapacity
|
||
let keys : Array Key := .mkEmpty initCapacity
|
||
mkPathAux (root := true) (todo.push e) keys config
|
||
|
||
/- Monad for finding matches while resolving deferred patterns. -/
|
||
@[reducible]
|
||
private def MatchM α := ReaderT WhnfCoreConfig (StateRefT (Array (Trie α)) MetaM)
|
||
|
||
private def runMatch (d : LazyDiscrTree α) (m : MatchM α β) : MetaM (β × LazyDiscrTree α) := do
|
||
let { config := c, tries := a, roots := r } := d
|
||
let (result, a) ← withReducible $ (m.run c).run a
|
||
pure (result, { config := c, tries := a, roots := r})
|
||
|
||
private def setTrie (i : TrieIndex) (v : Trie α) : MatchM α Unit :=
|
||
modify (·.set! i v)
|
||
|
||
/-- Create a new trie with the given lazy entry. -/
|
||
private def newTrie [Monad m] [MonadState (Array (Trie α)) m] (e : LazyEntry α) : m TrieIndex := do
|
||
modifyGet fun a => let sz := a.size; (sz, a.push (.node #[] 0 {} #[e]))
|
||
|
||
/-- Add a lazy entry to an existing trie. -/
|
||
private def addLazyEntryToTrie (i:TrieIndex) (e : LazyEntry α) : MatchM α Unit :=
|
||
modify (·.modify i (·.pushPending e))
|
||
|
||
/--
|
||
This evaluates all lazy entries in a trie and updates `values`, `starIdx`, and `children`
|
||
accordingly.
|
||
-/
|
||
private partial def evalLazyEntries (config : WhnfCoreConfig)
|
||
(values : Array α) (starIdx : TrieIndex) (children : HashMap Key TrieIndex)
|
||
(entries : Array (LazyEntry α)) :
|
||
MatchM α (Array α × TrieIndex × HashMap Key TrieIndex) := do
|
||
let rec iter values starIdx children (i : Nat) : MatchM α _ := do
|
||
if p : i < entries.size then
|
||
let (todo, lctx, v) := entries[i]
|
||
if todo.isEmpty then
|
||
let values := values.push v
|
||
iter values starIdx children (i+1)
|
||
else
|
||
let e := todo.back
|
||
let todo := todo.pop
|
||
let (k, todo) ← withLCtx lctx.1 lctx.2 $ pushArgs false todo e config
|
||
if k == .star then
|
||
if starIdx = 0 then
|
||
let starIdx ← newTrie (todo, lctx, v)
|
||
iter values starIdx children (i+1)
|
||
else
|
||
addLazyEntryToTrie starIdx (todo, lctx, v)
|
||
iter values starIdx children (i+1)
|
||
else
|
||
match children.find? k with
|
||
| none =>
|
||
let children := children.insert k (← newTrie (todo, lctx, v))
|
||
iter values starIdx children (i+1)
|
||
| some idx =>
|
||
addLazyEntryToTrie idx (todo, lctx, v)
|
||
iter values starIdx children (i+1)
|
||
else
|
||
pure (values, starIdx, children)
|
||
iter values starIdx children 0
|
||
|
||
private def evalNode (c : TrieIndex) :
|
||
MatchM α (Array α × TrieIndex × HashMap Key TrieIndex) := do
|
||
let .node vs star cs pending := (←get).get! c
|
||
if pending.size = 0 then
|
||
pure (vs, star, cs)
|
||
else
|
||
let config ← read
|
||
setTrie c default
|
||
let (vs, star, cs) ← evalLazyEntries config vs star cs pending
|
||
setTrie c <| .node vs star cs #[]
|
||
pure (vs, star, cs)
|
||
|
||
/--
|
||
Return the information about the trie at the given idnex.
|
||
|
||
Used for internal debugging purposes.
|
||
-/
|
||
private def getTrie (d : LazyDiscrTree α) (idx : TrieIndex) :
|
||
MetaM ((Array α × TrieIndex × HashMap Key TrieIndex) × LazyDiscrTree α) :=
|
||
runMatch d (evalNode idx)
|
||
|
||
/--
|
||
A match result contains the terms formed from matching a term against
|
||
patterns in the discrimination tree.
|
||
|
||
-/
|
||
private structure MatchResult (α : Type) where
|
||
/--
|
||
The elements in the match result.
|
||
|
||
The top-level array represents an array from `score` values to the
|
||
results with that score. A `score` is the number of non-star matches
|
||
in a pattern against the term, and thus bounded by the size of the
|
||
term being matched against. The elements of this array are themselves
|
||
arrays of non-empty arrays so that we can defer concatenating results until
|
||
needed.
|
||
-/
|
||
elts : Array (Array (Array α)) := #[]
|
||
|
||
private def MatchResult.push (r : MatchResult α) (score : Nat) (e : Array α) : MatchResult α :=
|
||
if e.isEmpty then
|
||
r
|
||
else if score < r.elts.size then
|
||
{ elts := r.elts.modify score (·.push e) }
|
||
else
|
||
let rec loop (a : Array (Array (Array α))) :=
|
||
if a.size < score then
|
||
loop (a.push #[])
|
||
else
|
||
{ elts := a.push #[e] }
|
||
termination_by score - a.size
|
||
loop r.elts
|
||
|
||
private partial def MatchResult.toArray (mr : MatchResult α) : Array α :=
|
||
loop (Array.mkEmpty n) mr.elts
|
||
where n := mr.elts.foldl (fun i a => a.foldl (fun n a => n + a.size) i) 0
|
||
loop (r : Array α) (a : Array (Array (Array α))) :=
|
||
if a.isEmpty then
|
||
r
|
||
else
|
||
loop (a.back.foldl (init := r) (fun r a => r ++ a)) a.pop
|
||
|
||
private partial def getMatchLoop (todo : Array Expr) (score : Nat) (c : TrieIndex)
|
||
(result : MatchResult α) : MatchM α (MatchResult α) := do
|
||
let (vs, star, cs) ← evalNode c
|
||
if todo.isEmpty then
|
||
return result.push score vs
|
||
else if star == 0 && cs.isEmpty then
|
||
return result
|
||
else
|
||
let e := todo.back
|
||
let todo := todo.pop
|
||
/- We must always visit `Key.star` edges since they are wildcards.
|
||
Thus, `todo` is not used linearly when there is `Key.star` edge
|
||
and there is an edge for `k` and `k != Key.star`. -/
|
||
let visitStar (result : MatchResult α) : MatchM α (MatchResult α) :=
|
||
if star != 0 then
|
||
getMatchLoop todo score star result
|
||
else
|
||
return result
|
||
let visitNonStar (k : Key) (args : Array Expr) (result : MatchResult α) :=
|
||
match cs.find? k with
|
||
| none => return result
|
||
| some c => getMatchLoop (todo ++ args) (score + 1) c result
|
||
let result ← visitStar result
|
||
let (k, args) ← MatchClone.getMatchKeyArgs e (root := false) (←read)
|
||
match k with
|
||
| .star => return result
|
||
/-
|
||
Note: dep-arrow vs arrow
|
||
Recall that dependent arrows are `(Key.other, #[])`, and non-dependent arrows are
|
||
`(Key.arrow, #[a, b])`.
|
||
A non-dependent arrow may be an instance of a dependent arrow (stored at `DiscrTree`).
|
||
Thus, we also visit the `Key.other` child.
|
||
-/
|
||
| .arrow => visitNonStar .other #[] (← visitNonStar k args result)
|
||
| _ => visitNonStar k args result
|
||
|
||
private def getStarResult (root : Lean.HashMap Key TrieIndex) : MatchM α (MatchResult α) :=
|
||
match root.find? .star with
|
||
| none =>
|
||
pure <| {}
|
||
| some idx => do
|
||
let (vs, _) ← evalNode idx
|
||
pure <| ({} : MatchResult α).push 0 vs
|
||
|
||
private def getMatchRoot (r : Lean.HashMap Key TrieIndex) (k : Key) (args : Array Expr)
|
||
(result : MatchResult α) : MatchM α (MatchResult α) :=
|
||
match r.find? k with
|
||
| none => pure result
|
||
| some c => getMatchLoop args 1 c result
|
||
|
||
/--
|
||
Find values that match `e` in `root`.
|
||
-/
|
||
private def getMatchCore (root : Lean.HashMap Key TrieIndex) (e : Expr) :
|
||
MatchM α (MatchResult α) := do
|
||
let result ← getStarResult root
|
||
let (k, args) ← MatchClone.getMatchKeyArgs e (root := true) (←read)
|
||
match k with
|
||
| .star => return result
|
||
/- See note about "dep-arrow vs arrow" at `getMatchLoop` -/
|
||
| .arrow =>
|
||
getMatchRoot root k args (←getMatchRoot root .other #[] result)
|
||
| _ =>
|
||
getMatchRoot root k args result
|
||
|
||
/--
|
||
Find values that match `e` in `d`.
|
||
|
||
The results are ordered so that the longest matches in terms of number of
|
||
non-star keys are first with ties going to earlier operators first.
|
||
-/
|
||
def getMatch (d : LazyDiscrTree α) (e : Expr) : MetaM (Array α × LazyDiscrTree α) :=
|
||
withReducible <| runMatch d <| (·.toArray) <$> getMatchCore d.roots e
|
||
|
||
/--
|
||
Structure for quickly initializing a lazy discrimination tree with a large number
|
||
of elements using concurrent functions for generating entries.
|
||
-/
|
||
private structure PreDiscrTree (α : Type) where
|
||
/-- Maps keys to index in tries array. -/
|
||
roots : HashMap Key Nat := {}
|
||
/-- Lazy entries for root of trie. -/
|
||
tries : Array (Array (LazyEntry α)) := #[]
|
||
deriving Inhabited
|
||
|
||
namespace PreDiscrTree
|
||
|
||
private def modifyAt (d : PreDiscrTree α) (k : Key)
|
||
(f : Array (LazyEntry α) → Array (LazyEntry α)) : PreDiscrTree α :=
|
||
let { roots, tries } := d
|
||
match roots.find? k with
|
||
| .none =>
|
||
let roots := roots.insert k tries.size
|
||
{ roots, tries := tries.push (f #[]) }
|
||
| .some i =>
|
||
{ roots, tries := tries.modify i f }
|
||
|
||
/-- Add an entry to the pre-discrimination tree.-/
|
||
private def push (d : PreDiscrTree α) (k : Key) (e : LazyEntry α) : PreDiscrTree α :=
|
||
d.modifyAt k (·.push e)
|
||
|
||
/-- Convert a pre-discrimination tree to a lazy discrimination tree. -/
|
||
private def toLazy (d : PreDiscrTree α) (config : WhnfCoreConfig := {}) : LazyDiscrTree α :=
|
||
let { roots, tries } := d
|
||
{ config, roots, tries := tries.map (.node {} 0 {}) }
|
||
|
||
/-- Merge two discrimination trees. -/
|
||
protected def append (x y : PreDiscrTree α) : PreDiscrTree α :=
|
||
let (x, y, f) :=
|
||
if x.roots.size ≥ y.roots.size then
|
||
(x, y, fun y x => x ++ y)
|
||
else
|
||
(y, x, fun x y => x ++ y)
|
||
let { roots := yk, tries := ya } := y
|
||
yk.fold (init := x) fun d k yi => d.modifyAt k (f ya[yi]!)
|
||
|
||
instance : Append (PreDiscrTree α) where
|
||
append := PreDiscrTree.append
|
||
|
||
end PreDiscrTree
|
||
|
||
/-- Initial entry in lazy discrimination tree -/
|
||
@[reducible]
|
||
structure InitEntry (α : Type) where
|
||
/-- Return root key for an entry. -/
|
||
key : Key
|
||
/-- Returns rest of entry for later insertion. -/
|
||
entry : LazyEntry α
|
||
|
||
namespace InitEntry
|
||
|
||
/--
|
||
Constructs an initial entry from an expression and value.
|
||
-/
|
||
def fromExpr (expr : Expr) (value : α) (config : WhnfCoreConfig := {}) : MetaM (InitEntry α) := do
|
||
let lctx ← getLCtx
|
||
let linst ← getLocalInstances
|
||
let lctx := (lctx, linst)
|
||
let (key, todo) ← LazyDiscrTree.rootKey config expr
|
||
pure <| { key, entry := (todo, lctx, value) }
|
||
|
||
/--
|
||
Creates an entry for a subterm of an initial entry.
|
||
|
||
This is slightly more efficient than using `fromExpr` on subterms since it avoids a redundant call
|
||
to `whnf`.
|
||
-/
|
||
def mkSubEntry (e : InitEntry α) (idx : Nat) (value : α) (config : WhnfCoreConfig := {}) :
|
||
MetaM (InitEntry α) := do
|
||
let (todo, lctx, _) := e.entry
|
||
let (key, todo) ← LazyDiscrTree.rootKey config todo[idx]!
|
||
pure <| { key, entry := (todo, lctx, value) }
|
||
|
||
end InitEntry
|
||
|
||
/-- Information about a failed import. -/
|
||
private structure ImportFailure where
|
||
/-- Module with constant that import failed on. -/
|
||
module : Name
|
||
/-- Constant that import failed on. -/
|
||
const : Name
|
||
/-- Exception that triggers error. -/
|
||
exception : Exception
|
||
|
||
/-- Information generation from imported modules. -/
|
||
private structure ImportData where
|
||
cache : IO.Ref (Lean.Meta.Cache)
|
||
errors : IO.Ref (Array ImportFailure)
|
||
|
||
private def ImportData.new : BaseIO ImportData := do
|
||
let cache ← IO.mkRef {}
|
||
let errors ← IO.mkRef #[]
|
||
pure { cache, errors }
|
||
|
||
private def addConstImportData
|
||
(env : Environment)
|
||
(modName : Name)
|
||
(d : ImportData)
|
||
(tree : PreDiscrTree α)
|
||
(act : Name → ConstantInfo → MetaM (Array (InitEntry α)))
|
||
(name : Name) (constInfo : ConstantInfo) : BaseIO (PreDiscrTree α) := do
|
||
if constInfo.isUnsafe then return tree
|
||
if !allowCompletion env name then return tree
|
||
let mstate : Meta.State := { cache := ←d.cache.get }
|
||
d.cache.set {}
|
||
let ctx : Meta.Context := { config := { transparency := .reducible } }
|
||
let cm := (act name constInfo).run ctx mstate
|
||
let cctx : Core.Context := {
|
||
fileName := default,
|
||
fileMap := default
|
||
}
|
||
let cstate : Core.State := {env}
|
||
match ←(cm.run cctx cstate).toBaseIO with
|
||
| .ok ((a, ms), _) =>
|
||
d.cache.set ms.cache
|
||
pure <| a.foldl (fun t e => t.push e.key e.entry) tree
|
||
| .error e =>
|
||
let i : ImportFailure := {
|
||
module := modName,
|
||
const := name,
|
||
exception := e
|
||
}
|
||
d.errors.modify (·.push i)
|
||
pure tree
|
||
|
||
/--
|
||
Contains the pre discrimination tree and any errors occuring during initialization of
|
||
the library search tree.
|
||
-/
|
||
private structure InitResults (α : Type) where
|
||
tree : PreDiscrTree α := {}
|
||
errors : Array ImportFailure := #[]
|
||
|
||
instance : Inhabited (InitResults α) where
|
||
default := {}
|
||
|
||
namespace InitResults
|
||
|
||
/-- Combine two initial results. -/
|
||
protected def append (x y : InitResults α) : InitResults α :=
|
||
let { tree := xv, errors := xe } := x
|
||
let { tree := yv, errors := ye } := y
|
||
{ tree := xv ++ yv, errors := xe ++ ye }
|
||
|
||
instance : Append (InitResults α) where
|
||
append := InitResults.append
|
||
|
||
end InitResults
|
||
|
||
private def toFlat (d : ImportData) (tree : PreDiscrTree α) :
|
||
BaseIO (InitResults α) := do
|
||
let de ← d.errors.swap #[]
|
||
pure ⟨tree, de⟩
|
||
|
||
private partial def loadImportedModule (env : Environment)
|
||
(act : Name → ConstantInfo → MetaM (Array (InitEntry α)))
|
||
(d : ImportData)
|
||
(tree : PreDiscrTree α)
|
||
(mname : Name)
|
||
(mdata : ModuleData)
|
||
(i : Nat := 0) : BaseIO (PreDiscrTree α) := do
|
||
if h : i < mdata.constNames.size then
|
||
let name := mdata.constNames[i]
|
||
let constInfo := mdata.constants[i]!
|
||
let tree ← addConstImportData env mname d tree act name constInfo
|
||
loadImportedModule env act d tree mname mdata (i+1)
|
||
else
|
||
pure tree
|
||
|
||
private def createImportedEnvironmentSeq (env : Environment)
|
||
(act : Name → ConstantInfo → MetaM (Array (InitEntry α)))
|
||
(start stop : Nat) : BaseIO (InitResults α) :=
|
||
do go (← ImportData.new) {} start stop
|
||
where go d (tree : PreDiscrTree α) (start stop : Nat) : BaseIO _ := do
|
||
if start < stop then
|
||
let mname := env.header.moduleNames[start]!
|
||
let mdata := env.header.moduleData[start]!
|
||
let tree ← loadImportedModule env act d tree mname mdata
|
||
go d tree (start+1) stop
|
||
else
|
||
toFlat d tree
|
||
termination_by stop - start
|
||
|
||
/-- Get the results of each task and merge using combining function -/
|
||
private def combineGet [Append α] (z : α) (tasks : Array (Task α)) : α :=
|
||
tasks.foldl (fun x t => x ++ t.get) (init := z)
|
||
|
||
/-- Create an imported environment for tree. -/
|
||
def createImportedEnvironment (env : Environment)
|
||
(act : Name → ConstantInfo → MetaM (Array (InitEntry α)))
|
||
(constantsPerTask : Nat := 1000) :
|
||
EIO Exception (LazyDiscrTree α) := do
|
||
let n := env.header.moduleData.size
|
||
let rec
|
||
/-- Allocate constants to tasks according to `constantsPerTask`. -/
|
||
go tasks start cnt idx := do
|
||
if h : idx < env.header.moduleData.size then
|
||
let mdata := env.header.moduleData[idx]
|
||
let cnt := cnt + mdata.constants.size
|
||
if cnt > constantsPerTask then
|
||
let t ← createImportedEnvironmentSeq env act start (idx+1) |>.asTask
|
||
go (tasks.push t) (idx+1) 0 (idx+1)
|
||
else
|
||
go tasks start cnt (idx+1)
|
||
else
|
||
if start < n then
|
||
tasks.push <$> (createImportedEnvironmentSeq env act start n).asTask
|
||
else
|
||
pure tasks
|
||
termination_by env.header.moduleData.size - idx
|
||
let tasks ← go #[] 0 0 0
|
||
let r := combineGet default tasks
|
||
if p : r.errors.size > 0 then
|
||
throw r.errors[0].exception
|
||
pure <| r.tree.toLazy
|