diff --git a/src/Lean/Meta/DiscrTree/Basic.lean b/src/Lean/Meta/DiscrTree/Basic.lean index 745033421224..e400cb2b0117 100644 --- a/src/Lean/Meta/DiscrTree/Basic.lean +++ b/src/Lean/Meta/DiscrTree/Basic.lean @@ -60,6 +60,8 @@ def Key.format : Key → Format instance : ToFormat Key := ⟨Key.format⟩ partial def Trie.format [ToFormat α] : Trie α → Format + | .chain k c => Format.group $ Format.paren $ + "chain " ++ Std.format k ++ " => " ++ format c | .node vs cs => Format.group $ Format.paren $ "node" ++ (if vs.isEmpty then Format.nil else " " ++ Std.format vs) ++ Format.join (cs.toList.map fun ⟨k, c⟩ => Format.line ++ Format.paren (Std.format k ++ " => " ++ format c)) @@ -126,7 +128,7 @@ private partial def createNodes (keys : Array Key) (v : α) (i : Nat) : Trie α if h : i < keys.size then let k := keys[i] let c := createNodes keys v (i+1) - .node #[] #[(k, c)] + .chain k c else .node #[v] #[] @@ -150,6 +152,16 @@ where termination_by vs.size - i private partial def insertAux [BEq α] (keys : Array Key) (v : α) : Nat → Trie α → Trie α + | i, .chain k c => + if h : i < keys.size then + if keys[i] == k then + .chain k (insertAux keys v (i+1) c) + else if keys[i] < k then + .node #[] #[(keys[i], createNodes keys v (i+1)), (k, c)] + else + .node #[] #[(k, c), (keys[i], createNodes keys v (i+1))] + else + .node #[v] #[(k, c)] | i, .node vs cs => if h : i < keys.size then let k := keys[i] diff --git a/src/Lean/Meta/DiscrTree/Main.lean b/src/Lean/Meta/DiscrTree/Main.lean index bdfbcdff9d7a..b34767cabcc9 100644 --- a/src/Lean/Meta/DiscrTree/Main.lean +++ b/src/Lean/Meta/DiscrTree/Main.lean @@ -430,6 +430,7 @@ private def getStarResult (d : DiscrTree α) : Array α := let result : Array α := .mkEmpty initCapacity match d.root.find? .star with | none => result + | some (.chain _ _) => result -- unreachable in well-formed trees! | some (.node vs _) => result ++ vs private abbrev findKey (cs : Array (Key × Trie α)) (k : Key) : Option (Key × Trie α) := @@ -437,6 +438,21 @@ private abbrev findKey (cs : Array (Key × Trie α)) (k : Key) : Option (Key × private partial def getMatchLoop (todo : Array Expr) (c : Trie α) (result : Array α) : MetaM (Array α) := do match c with + | .chain key child => + if todo.isEmpty then + return result + else + let e := todo.back! + let todo := todo.pop + let (k, args) ← getMatchKeyArgs e (root := false) + if key == .star then + getMatchLoop todo child result + else + if key == k then + getMatchLoop (todo ++ args) child result + else + return result + | .node vs cs => if todo.isEmpty then return result ++ vs @@ -450,19 +466,17 @@ private partial def getMatchLoop (todo : Array Expr) (c : Trie α) (result : Arr /- 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 : Array α) : MetaM (Array α) := + let result ← if first.1 == .star then getMatchLoop todo first.2 result else - return result - let visitNonStar (k : Key) (args : Array Expr) (result : Array α) : MetaM (Array α) := + pure result + match k with + | .star => return result + | _ => match findKey cs k with | none => return result | some c => getMatchLoop (todo ++ args) c.2 result - let result ← visitStar result - match k with - | .star => return result - | _ => visitNonStar k args result private def getMatchRoot (d : DiscrTree α) (k : Key) (args : Array Expr) (result : Array α) : MetaM (Array α) := match d.root.find? k with @@ -544,6 +558,7 @@ private partial def getAllValuesForKey (d : DiscrTree α) (k : Key) (result : Ar where go (trie : Trie α) (result : Array α) : Array α := Id.run do match trie with + | .chain _ c => go c result | .node vs cs => let mut result := result ++ vs for (_, trie) in cs do @@ -577,11 +592,28 @@ partial def getUnify (d : DiscrTree α) (e : Expr) : MetaM (Array α) := where process (skip : Nat) (todo : Array Expr) (c : Trie α) (result : Array α) : MetaM (Array α) := do match skip, c with + | skip+1, .chain key child => + process (skip + key.arity) todo child result | skip+1, .node _ cs => if cs.isEmpty then return result else cs.foldlM (init := result) fun result ⟨k, c⟩ => process (skip + k.arity) todo c result + | 0, .chain key child => + if todo.isEmpty then + return result + else + let e := todo.back! + let todo := todo.pop + let (k, args) ← getUnifyKeyArgs e (root := false) + if k == .star then + process key.arity todo child result + else if key == .star then + process 0 todo child result + else if key == k then + process 0 (todo ++ args) child result + else + return result | 0, .node vs cs => do if todo.isEmpty then return result ++ vs @@ -590,19 +622,18 @@ where else let e := todo.back! let todo := todo.pop + let first := cs[0]! let (k, args) ← getUnifyKeyArgs e (root := false) - let visitStar (result : Array α) : MetaM (Array α) := - let first := cs[0]! - if first.1 == .star then - process 0 todo first.2 result - else - return result - let visitNonStar (k : Key) (args : Array Expr) (result : Array α) : MetaM (Array α) := + match k with + | .star => cs.foldlM (init := result) fun result ⟨k, c⟩ => process k.arity todo c result + | _ => + let result ← + if first.1 == .star then + process 0 todo first.2 result + else + pure result match findKey cs k with - | none => return result + | none => return result | some c => process 0 (todo ++ args) c.2 result - match k with - | .star => cs.foldlM (init := result) fun result ⟨k, c⟩ => process k.arity todo c result - | _ => visitNonStar k args (← visitStar result) end Lean.Meta.DiscrTree diff --git a/src/Lean/Meta/DiscrTree/Types.lean b/src/Lean/Meta/DiscrTree/Types.lean index 45e29c6b0b24..9af3726b6065 100644 --- a/src/Lean/Meta/DiscrTree/Types.lean +++ b/src/Lean/Meta/DiscrTree/Types.lean @@ -38,6 +38,12 @@ instance : Hashable Key := ⟨Key.hash⟩ Discrimination tree trie. See `DiscrTree`. -/ inductive Trie (α : Type) where + /-- + In many uses cases, tries do a lot of their actual discrimination near the root and end up with + long chains of singleton nodes. The special-cased `.chain key child` node is equivalent to + `.node #[] #[(key, child)]` but is more memory efficient. + -/ + | chain (key : Key) (child : Trie α) : Trie α | node (vs : Array α) (children : Array (Key × Trie α)) : Trie α end DiscrTree diff --git a/src/Lean/Meta/DiscrTree/Util.lean b/src/Lean/Meta/DiscrTree/Util.lean index bd95cbfc7162..8e5100aedbb9 100644 --- a/src/Lean/Meta/DiscrTree/Util.lean +++ b/src/Lean/Meta/DiscrTree/Util.lean @@ -15,6 +15,8 @@ Monadically fold the keys and values stored in a `Trie`. @[specialize] partial def foldM [Monad m] (initialKeys : Array Key) (f : σ → Array Key → α → m σ) : (init : σ) → Trie α → m σ + | init, Trie.chain k c => + c.foldM (initialKeys.push k) f init | init, Trie.node vs children => do let s ← vs.foldlM (init := init) fun s v => f s initialKeys v children.foldlM (init := s) fun s (k, t) => @@ -32,6 +34,8 @@ Monadically fold the values stored in a `Trie`. -/ @[specialize] partial def foldValuesM [Monad m] (f : σ → α → m σ) : (init : σ) → Trie α → m σ + | init, chain _ c => + c.foldValuesM (init := init) f | init, node vs children => do let s ← vs.foldlM (init := init) f children.foldlM (init := s) fun s (_, c) => c.foldValuesM (init := s) f @@ -47,6 +51,7 @@ def foldValues (f : σ → α → σ) (init : σ) (t : Trie α) : σ := The number of values stored in a `Trie`. -/ partial def size : Trie α → Nat + | Trie.chain _ c => size c | Trie.node vs children => children.foldl (init := vs.size) fun n (_, c) => n + size c @@ -55,13 +60,17 @@ Generate a trie node from values and an array of children. -/ @[inline] def mkNode (vs : Array α) (cs : Array (Key × Trie α)) : Trie α := - .node vs cs + if h : vs.isEmpty ∧ cs.size = 1 then + .chain cs[0].1 cs[0].2 + else + .node vs cs /-- Inspect a trie node as an array of values and an array of children. -/ @[inline] def asNode : Trie α → Array α × Array (Key × Trie α) + | .chain k v => ⟨#[], #[(k, v)]⟩ | .node vs cs => ⟨vs, cs⟩ /-- @@ -70,6 +79,7 @@ Equivalent to `t.asNode.1`. -/ @[inline] def nodeValues : Trie α → Array α + | .chain _ _ => #[] | .node vs _ => vs /-- @@ -78,6 +88,7 @@ Equivalent to `t.asNode.2`. -/ @[inline] def nodeChildren : Trie α → Array (Key × Trie α) + | .chain k v => #[(k, v)] | .node _ cs => cs /-- @@ -88,6 +99,7 @@ invariant that no trie node has an empty child node. -/ @[inline] def isEmptyNode : Trie α → Bool + | .chain _ _ => false | .node vs children => vs.isEmpty && children.isEmpty end Trie @@ -161,6 +173,11 @@ Any resulting subtrees containing no values will be pruned. partial def Trie.mapArraysM (t : DiscrTree.Trie α) (f : Array α → m (Array β)) : m (DiscrTree.Trie β) := match t with + | .chain k c => do + let vs ← f #[] -- Corner case. Possible future optimization: modify `mapArraysM` semantics to only call `f` for non-empty arrays and then eliminate this line + let c ← c.mapArraysM f + let cs := if c.isEmptyNode then #[] else #[(k, c)] + return Trie.mkNode vs cs | .node vs children => do let vs ← f vs let children ← children.filterMapM fun (k, child) => do @@ -169,7 +186,7 @@ partial def Trie.mapArraysM (t : DiscrTree.Trie α) (f : Array α → m (Array return none else return some (k, child) - return .node vs children + return Trie.mkNode vs children /-- Apply a monadic function to the array of values at each node in a `DiscrTree`. -/ @[inline] diff --git a/src/Lean/Meta/Sym/Simp/DiscrTree.lean b/src/Lean/Meta/Sym/Simp/DiscrTree.lean index 2152e1a9a0c9..fb59a4181e58 100644 --- a/src/Lean/Meta/Sym/Simp/DiscrTree.lean +++ b/src/Lean/Meta/Sym/Simp/DiscrTree.lean @@ -167,6 +167,19 @@ def pushArgsTodo (todo : Array Expr) (e : Expr) : Array Expr := partial def getMatchLoop (mctx : MetavarContext) (todo : Array Expr) (c : Trie α) (result : Array α) : Array α := match c with + | .chain key child => + if todo.isEmpty then + result + else + let e := resolveAssignedMVars mctx <| etaReduce todo.back! + let todo := todo.pop + if key == .star then + getMatchLoop mctx todo child result + else if key == getKey e then + getMatchLoop mctx (pushArgsTodo todo e) child result + else + result + | .node vs cs => let csize := cs.size if todo.isEmpty then @@ -177,25 +190,17 @@ partial def getMatchLoop (mctx : MetavarContext) (todo : Array Expr) (c : Trie let e := resolveAssignedMVars mctx <| etaReduce todo.back! let todo := todo.pop let first := cs[0] /- Recall that `Key.star` is the minimal key -/ - if csize = 1 then - /- Special case: only one child node -/ + /- 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 result := if first.1 == .star then getMatchLoop mctx todo first.2 result - else if first.1 == getKey e then - getMatchLoop mctx (pushArgsTodo todo e) first.2 result - else - result - else - /- 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 result := if first.1 == .star then - getMatchLoop mctx todo first.2 result else result - match findKey? cs (getKey e) with - | none => result - | some c => getMatchLoop mctx (pushArgsTodo todo e) c.2 result + match findKey? cs (getKey e) with + | none => result + | some c => getMatchLoop mctx (pushArgsTodo todo e) c.2 result /-- Retrieves all values whose patterns match the expression `e`. @@ -204,6 +209,7 @@ Retrieves all values whose patterns match the expression `e`. public def getMatch (mctx : MetavarContext) (d : DiscrTree α) (e : Expr) : Array α := let result := match d.root.find? .star with | none => .mkEmpty initCapacity + | some (.chain _ _) => .mkEmpty initCapacity -- unreachable in well-formed trees! | some (.node vs _) => vs let e := resolveAssignedMVars mctx <| etaReduce e match d.root.find? (getKey e) with diff --git a/tests/elab/discrTreeFind.lean b/tests/elab/discrTreeFind.lean index b33afa78449b..e4e63451be96 100644 --- a/tests/elab/discrTreeFind.lean +++ b/tests/elab/discrTreeFind.lean @@ -13,13 +13,13 @@ opaque h : Nat → Nat → Nat → Nat → Nat /-- info: [0] -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))))) --- info: [0, 2] -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))) (a => (node #[2])))) --- info: [0, 2] -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])) (b => (node #[3])))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))) (a => (node #[2])) (b => (node #[3])))) --- info: [0, 2] $(f => (node @@ -46,13 +46,13 @@ info: [0, 1, 4] /-- info: [0] -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))))) --- info: [0] -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))) (a => (node #[2])))) --- info: [0] -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])) (b => (node #[3])))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))) (a => (node #[2])) (b => (node #[3])))) --- info: [0] $(f => (node @@ -79,13 +79,13 @@ info: [0, 1] /-- info: [0] -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))))) --- info: [0, 2] -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))) (a => (node #[2])))) --- info: [0, 2] -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])) (b => (node #[3])))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))) (a => (node #[2])) (b => (node #[3])))) --- info: [0, 2] $(f => (node @@ -112,13 +112,13 @@ info: [0, 1, 4] /-- info: [0, 1] -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))))) --- info: [0, 1, 2] -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))) (a => (node #[2])))) --- info: [0, 1, 2, 3] -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])) (b => (node #[3])))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))) (a => (node #[2])) (b => (node #[3])))) --- info: [0, 1, 4, 2, 3] $(f => (node @@ -145,13 +145,13 @@ info: [0, 1, 4] /-- info: ([0, 1], 1) -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))))) --- info: ([0, 1, 2], 1) -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))) (a => (node #[2])))) --- info: ([0, 1, 2, 3], 1) -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])) (b => (node #[3])))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))) (a => (node #[2])) (b => (node #[3])))) --- info: ([0, 1, 4, 2, 3], 1) $(f => (node @@ -178,13 +178,13 @@ info: ([0, 1, 4, 2, 3], 1) /-- info: ([0, 1], 1) -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))))) --- info: ([0, 1, 2], 1) -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))) (a => (node #[2])))) --- info: ([0, 1, 2, 3], 1) -$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])) (b => (node #[3])))) +$(f => (node (* => (node #[0])) (f => (chain * => (node #[1]))) (a => (node #[2])) (b => (node #[3])))) --- info: ([0, 1, 4, 2, 3], 1) $(f => (node diff --git a/tests/elab/discrTreeOps.lean b/tests/elab/discrTreeOps.lean index 5165da64a8f1..c9360f81e818 100644 --- a/tests/elab/discrTreeOps.lean +++ b/tests/elab/discrTreeOps.lean @@ -10,23 +10,23 @@ opaque h : Nat → Nat → Nat /-- info: 1 | [([f, 1], 1)] -$(f => (node (1 => (node #[1])))) +$(f => (chain 1 => (node #[1]))) --- info: 2 | [([f, 1], 1), ([f, 1], 2)] -$(f => (node (1 => (node #[1, 2])))) +$(f => (chain 1 => (node #[1, 2]))) --- info: 3 | [([f, 1], 1), ([f, 1], 2), ([f, 2], 3)] $(f => (node (1 => (node #[1, 2])) (2 => (node #[3])))) --- info: 4 | [([f, 1], 1), ([f, 1], 2), ([f, 2], 3), ([f, g, "a"], 4)] -$(f => (node (1 => (node #[1, 2])) (2 => (node #[3])) (g => (node ("a" => (node #[4])))))) +$(f => (node (1 => (node #[1, 2])) (2 => (node #[3])) (g => (chain "a" => (node #[4]))))) --- info: 5 | [([f, 1], 1), ([f, 1], 2), ([f, 2], 3), ([f, g, "a"], 4), ([f, h, 1, 2], 5)] $(f => (node (1 => (node #[1, 2])) (2 => (node #[3])) - (g => (node ("a" => (node #[4])))) - (h => (node (1 => (node (2 => (node #[5])))))))) + (g => (chain "a" => (node #[4]))) + (h => (chain 1 => (chain 2 => (node #[5])))))) -/ #guard_msgs in #eval do @@ -43,11 +43,11 @@ $(f => (node logInfo m!"{t.size} | {t.toArray}\n${t}" /-- -info: (f => (node (10 => (node #[1, 2])))) +info: (f => (chain 10 => (node #[1, 2]))) [([f, 10], 1), ([f, 10], 2)] [1, 2] true true false --- -info: (f => (node (10 => (node #[2, 1])))) +info: (f => (chain 10 => (node #[2, 1]))) [([f, 10], 2), ([f, 10], 1)] [2, 1] true true false -/ @@ -67,10 +67,9 @@ info: (f => (node (10 => (node #[2, 1])))) /-- -info: (f => (node - (h => (node - (0 => (node (0 => (node #[11])) (1 => (node #[12])))) - (1 => (node (0 => (node #[13])) (1 => (node #[14])))))))) +info: (f => (chain h => (node + (0 => (node (0 => (node #[11])) (1 => (node #[12])))) + (1 => (node (0 => (node #[13])) (1 => (node #[14]))))))) -/ #guard_msgs in #eval do @@ -82,36 +81,31 @@ info: (f => (node logInfo m!"{t.mapArrays (·.map (· + 10))}" /-- -info: -("A" => (node #[10, 11])) +info: ("A" => (node #[10, 11])) ("B" => (node #[12])) -(g => (node (0 => (node #[13, 14, 15])))) +(g => (chain 0 => (node #[13, 14, 15]))) (f => (node (0 => (node #[16])) (1 => (node #[17, 18])) (2 => (node #[19])) (f => (node (1 => (node #[20])) (2 => (node #[21])))))) --- -info: -("A" => (node #[0, 1])) +info: ("A" => (node #[0, 1])) ("B" => (node #[2])) -(g => (node (0 => (node #[3, 4, 5])))) -(f => (node (0 => (node #[6])) (1 => (node #[7, 8])) (2 => (node #[9])) (f => (node (1 => (node #[10])))))) +(g => (chain 0 => (node #[3, 4, 5]))) +(f => (node (0 => (node #[6])) (1 => (node #[7, 8])) (2 => (node #[9])) (f => (chain 1 => (node #[10]))))) --- -info: -("A" => (node #[0])) +info: ("A" => (node #[0])) ("B" => (node #[2])) -(g => (node (0 => (node #[4])))) -(f => (node (0 => (node #[6])) (1 => (node #[8])) (f => (node (1 => (node #[10])))))) +(g => (chain 0 => (node #[4]))) +(f => (node (0 => (node #[6])) (1 => (node #[8])) (f => (chain 1 => (node #[10]))))) --- -info: -(g => (node (0 => (node #[3, 4, 5])))) +info: (g => (chain 0 => (node #[3, 4, 5]))) --- -info: -("B" => (node #[2])) +info: ("B" => (node #[2])) (f => (node (0 => (node #[6])) (2 => (node #[9])) (f => (node (1 => (node #[10])) (2 => (node #[11])))))) --- -info: ("A" => (node #[0, 1])) (g => (node (0 => (node #[3, 4, 5])))) (f => (node (1 => (node #[7, 8])))) +info: ("A" => (node #[0, 1])) (g => (chain 0 => (node #[3, 4, 5]))) (f => (chain 1 => (node #[7, 8]))) --- info: -/ diff --git a/tests/elab/meta3.lean b/tests/elab/meta3.lean index f528c87a7268..a1b76fa402bf 100644 --- a/tests/elab/meta3.lean +++ b/tests/elab/meta3.lean @@ -53,11 +53,12 @@ do let d : DiscrTree Nat := {}; set_option trace.Meta.debug true in set_option pp.mvars false in /-- -trace: [Meta.debug] (Add.add => (node - (Nat => (node - (* => (node (* => (node (10 => (node #[1])) (20 => (node #[4])))) (0 => (node (10 => (node #[2])))))))))) +trace: [Meta.debug] +(Add.add => (chain Nat => (chain * => (node + (* => (node (10 => (node #[1])) (20 => (node #[4])))) + (0 => (chain 10 => (node #[2]))))))) (* => (node #[5])) - (Nat.add => (node (0 => (node (20 => (node #[3])))))) + (Nat.add => (chain 0 => (chain 20 => (node #[3])))) [Meta.debug] #[5, 1] [Meta.debug] Add.add ?_ ?_ [Meta.debug] #[5]