From 959bbe7e0b41cbcf679b3e1b98ac3a4fff37a715 Mon Sep 17 00:00:00 2001 From: Rob Simmons Date: Wed, 19 Aug 2026 12:15:53 -0400 Subject: [PATCH 1/6] fix: collapse empty DiscrTree tries --- src/Lean/Meta/DiscrTree/Util.lean | 25 +++- tests/elab/discrTreeFind.lean | 205 ++++++++++++++++++++++++++++++ tests/elab/discrTreeGrind.lean | 9 ++ tests/elab/discrTreeOps.lean | 112 ++++++++++++++++ 4 files changed, 348 insertions(+), 3 deletions(-) create mode 100644 tests/elab/discrTreeFind.lean create mode 100644 tests/elab/discrTreeGrind.lean create mode 100644 tests/elab/discrTreeOps.lean diff --git a/src/Lean/Meta/DiscrTree/Util.lean b/src/Lean/Meta/DiscrTree/Util.lean index 1916e78a59cd..1bc8f3c5d8d0 100644 --- a/src/Lean/Meta/DiscrTree/Util.lean +++ b/src/Lean/Meta/DiscrTree/Util.lean @@ -111,16 +111,35 @@ def size (t : DiscrTree α) : Nat := variable {m : Type → Type} [Monad m] +/-- +Checks that a trie node has no values and no children. + +This is only a check for actual trie emptiness (`t.size = 0`) if all operations maintain the +invariant that no trie node has an empty child node. +-/ +def Trie.isEmptyNode {α} : Trie α → Bool + | .node vs children => vs.isEmpty && children.isEmpty + /-- Apply a monadic function to the array of values at each node in a `DiscrTree`. -/ partial def Trie.mapArraysM (t : DiscrTree.Trie α) (f : Array α → m (Array β)) : m (DiscrTree.Trie β) := match t with - | .node vs children => - return .node (← f vs) (← children.mapM fun (k, t') => do pure (k, ← t'.mapArraysM f)) + | .node vs children => do + let vs ← f vs + let children ← children.filterMapM fun (k, child) => do + let child ← child.mapArraysM f + if child.isEmptyNode then + return none + else + return some (k, child) + return .node vs children /-- Apply a monadic function to the array of values at each node in a `DiscrTree`. -/ def mapArraysM (d : DiscrTree α) (f : Array α → m (Array β)) : m (DiscrTree β) := do - pure { root := ← d.root.mapM (fun t => t.mapArraysM f) } + let root ← d.root.mapM (fun t => t.mapArraysM f) + let emptyKeys := root.foldl (init := #[]) fun emptyKeys k t => + if t.isEmptyNode then emptyKeys.push k else emptyKeys + pure { root := emptyKeys.foldl (init := root) fun hashMap k => hashMap.erase k } /-- Apply a function to the array of values at each node in a `DiscrTree`. -/ def mapArrays (d : DiscrTree α) (f : Array α → Array β) : DiscrTree β := diff --git a/tests/elab/discrTreeFind.lean b/tests/elab/discrTreeFind.lean new file mode 100644 index 000000000000..b79340251b2d --- /dev/null +++ b/tests/elab/discrTreeFind.lean @@ -0,0 +1,205 @@ +import Lean +open Lean Meta + +opaque a : Nat +opaque b : Nat +opaque f : Nat → Nat +opaque h : Nat → Nat → Nat → Nat → Nat + +/-- +info: [0] +$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))))) +--- +info: [0, 2] +$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])))) +--- +info: [0, 2] +$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])) (b => (node #[3])))) +--- +info: [0, 2] +$(f => (node + (* => (node #[0])) + (f => (node (* => (node #[1])) (b => (node #[4])))) + (a => (node #[2])) + (b => (node #[3])))) +--- +info: [0, 1, 4] +-/ +#guard_msgs in +#eval do + let t : DiscrTree Nat := {} + let t ← t.insert (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat))) 0 + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))) 1 + logInfo m!"{← t.getMatch (mkApp (mkConst ``f) (mkConst ``a))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkConst ``a)) 2 + logInfo m!"{← t.getMatch (mkApp (mkConst ``f) (mkConst ``a))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkConst ``b)) 3 + logInfo m!"{← t.getMatch (mkApp (mkConst ``f) (mkConst ``a))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkConst ``f) (mkConst ``b))) 4 + logInfo m!"{← t.getMatch (mkApp (mkConst ``f) (mkConst ``a))}\n${t}" + logInfo m!"{← t.getMatch (mkApp (mkConst ``f) (mkApp (mkConst ``f) (mkConst ``b)))}" + +/-- +info: [0] +$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))))) +--- +info: [0] +$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])))) +--- +info: [0] +$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])) (b => (node #[3])))) +--- +info: [0] +$(f => (node + (* => (node #[0])) + (f => (node (* => (node #[1])) (b => (node #[4])))) + (a => (node #[2])) + (b => (node #[3])))) +--- +info: [0, 1] +-/ +#guard_msgs in +#eval do + let t : DiscrTree Nat := {} + let t ← t.insert (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat))) 0 + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))) 1 + logInfo m!"{← t.getMatch (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkConst ``a)) 2 + logInfo m!"{← t.getMatch (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkConst ``b)) 3 + logInfo m!"{← t.getMatch (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkConst ``f) (mkConst ``b))) 4 + logInfo m!"{← t.getMatch (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))}\n${t}" + logInfo m!"{← t.getMatch (mkApp (mkConst ``f) (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat))))}" + +/-- +info: [0] +$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))))) +--- +info: [0, 2] +$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])))) +--- +info: [0, 2] +$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])) (b => (node #[3])))) +--- +info: [0, 2] +$(f => (node + (* => (node #[0])) + (f => (node (* => (node #[1])) (b => (node #[4])))) + (a => (node #[2])) + (b => (node #[3])))) +--- +info: [0, 1, 4] +-/ +#guard_msgs in +#eval do + let t : DiscrTree Nat := {} + let t ← t.insert (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat))) 0 + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))) 1 + logInfo m!"{← t.getUnify (mkApp (mkConst ``f) (mkConst ``a))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkConst ``a)) 2 + logInfo m!"{← t.getUnify (mkApp (mkConst ``f) (mkConst ``a))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkConst ``b)) 3 + logInfo m!"{← t.getUnify (mkApp (mkConst ``f) (mkConst ``a))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkConst ``f) (mkConst ``b))) 4 + logInfo m!"{← t.getUnify (mkApp (mkConst ``f) (mkConst ``a))}\n${t}" + logInfo m!"{← t.getUnify (mkApp (mkConst ``f) (mkApp (mkConst ``f) (mkConst ``b)))}" + +/-- +info: [0, 1] +$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))))) +--- +info: [0, 1, 2] +$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])))) +--- +info: [0, 1, 2, 3] +$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))) (a => (node #[2])) (b => (node #[3])))) +--- +info: [0, 1, 4, 2, 3] +$(f => (node + (* => (node #[0])) + (f => (node (* => (node #[1])) (b => (node #[4])))) + (a => (node #[2])) + (b => (node #[3])))) +--- +info: [0, 1, 4] +-/ +#guard_msgs in +#eval do + let t : DiscrTree Nat := {} + let t ← t.insert (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat))) 0 + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))) 1 + logInfo m!"{← t.getUnify (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkConst ``a)) 2 + logInfo m!"{← t.getUnify (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkConst ``b)) 3 + logInfo m!"{← t.getUnify (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkConst ``f) (mkConst ``b))) 4 + logInfo m!"{← t.getUnify (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))}\n${t}" + logInfo m!"{← t.getUnify (mkApp (mkConst ``f) (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat))))}" + +/-- +info: ([0, 1], 1) +$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))))) +--- +info: ([0, 1, 2], 1) +$(f => (node (* => (node #[0])) (f => (node (* => (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])))) +--- +info: ([0, 1, 4, 2, 3], 1) +$(f => (node + (* => (node #[0])) + (f => (node (* => (node #[1])) (b => (node #[4])))) + (a => (node #[2])) + (b => (node #[3])))) +--- +info: ([0, 1, 4, 2, 3], 1) +-/ +#guard_msgs in +#eval do + let t : DiscrTree Nat := {} + let t ← t.insert (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat))) 0 + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))) 1 + logInfo m!"{← t.getMatchLiberal (mkApp (mkConst ``f) (mkConst ``a))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkConst ``a)) 2 + logInfo m!"{← t.getMatchLiberal (mkApp (mkConst ``f) (mkConst ``a))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkConst ``b)) 3 + logInfo m!"{← t.getMatchLiberal (mkApp (mkConst ``f) (mkConst ``a))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkConst ``f) (mkConst ``b))) 4 + logInfo m!"{← t.getMatchLiberal (mkApp (mkConst ``f) (mkConst ``a))}\n${t}" + logInfo m!"{← t.getMatchLiberal (mkApp (mkConst ``f) (mkApp (mkConst ``f) (mkConst ``b)))}" + +/-- +info: ([0, 1], 1) +$(f => (node (* => (node #[0])) (f => (node (* => (node #[1])))))) +--- +info: ([0, 1, 2], 1) +$(f => (node (* => (node #[0])) (f => (node (* => (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])))) +--- +info: ([0, 1, 4, 2, 3], 1) +$(f => (node + (* => (node #[0])) + (f => (node (* => (node #[1])) (b => (node #[4])))) + (a => (node #[2])) + (b => (node #[3])))) +--- +info: ([0, 1, 4, 2, 3], 1) +-/ +#guard_msgs in +#eval do + let t : DiscrTree Nat := {} + let t ← t.insert (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat))) 0 + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))) 1 + logInfo m!"{← t.getMatchLiberal (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkConst ``a)) 2 + logInfo m!"{← t.getMatchLiberal (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkConst ``b)) 3 + logInfo m!"{← t.getMatchLiberal (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkConst ``f) (mkConst ``b))) 4 + logInfo m!"{← t.getMatchLiberal (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat)))}\n${t}" + logInfo m!"{← t.getMatchLiberal (mkApp (mkConst ``f) (mkApp (mkConst ``f) (← mkFreshExprMVar (mkConst ``Nat))))}" diff --git a/tests/elab/discrTreeGrind.lean b/tests/elab/discrTreeGrind.lean new file mode 100644 index 000000000000..fa1abebbf241 --- /dev/null +++ b/tests/elab/discrTreeGrind.lean @@ -0,0 +1,9 @@ +def F (x : Nat) : Nat := x +def G (x : Nat) : Nat := x +def H (x : Nat) : Nat := x + +-- LHS `F (G (H x))` -> key path [F, G, H, *]: a 3-long chain under the root. +@[grind =] theorem FGH (x : Nat) : F (G (H x)) = x := rfl + +example (y : Nat) : F (G (H y)) = y := by grind +example (y : Nat) : F (G (H (F (G (H y))))) = y := by grind diff --git a/tests/elab/discrTreeOps.lean b/tests/elab/discrTreeOps.lean new file mode 100644 index 000000000000..5f1a23fb59ce --- /dev/null +++ b/tests/elab/discrTreeOps.lean @@ -0,0 +1,112 @@ +import Lean +open Lean Meta + +opaque f : Nat → Nat +opaque g : String → Nat +opaque h : Nat → Nat → Nat + +/-- +info: 1 | [([f, 1], 1)] +$(f => (node (1 => (node #[1])))) +--- +info: 2 | [([f, 1], 1), ([f, 1], 2)] +$(f => (node (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])))))) +--- +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])))))))) +-/ +#guard_msgs in +#eval do + let t : DiscrTree Nat := {} + let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 1)) 1 + logInfo m!"{t.size} | {t.toArray}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 1)) 2 + logInfo m!"{t.size} | {t.toArray}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 2)) 3 + logInfo m!"{t.size} | {t.toArray}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkConst ``g) (mkStrLit "a"))) 4 + logInfo m!"{t.size} | {t.toArray}\n${t}" + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkApp (mkConst ``h) (mkNatLit 1)) (mkNatLit 2))) 5 + logInfo m!"{t.size} | {t.toArray}\n${t}" + +/-- +info: (f => (node (10 => (node #[1, 2])))) +[([f, 10], 1), ([f, 10], 2)] +[1, 2] true true false +--- +info: (f => (node (10 => (node #[2, 1])))) +[([f, 10], 2), ([f, 10], 1)] +[2, 1] true true false +-/ +#guard_msgs in +#eval do + let t : DiscrTree Nat := {} + let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 10)) 1 + let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 10)) 2 + let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 10)) 2 + let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 10)) 1 + let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 10)) 2 + logInfo m!"{t}\n{t.toArray}\n{t.values} {t.containsValueP (· == 1)} {t.containsValueP (· == 2)} {t.containsValueP (· == 3)}" + let t : DiscrTree Nat := {} + let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 10)) 2 + let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 10)) 1 + logInfo m!"{t}\n{t.toArray}\n{t.values} {t.containsValueP (· == 1)} {t.containsValueP (· == 2)} {t.containsValueP (· == 3)}" + + +/-- +info: (f => (node + (h => (node + (0 => (node (0 => (node #[11])) (1 => (node #[12])))) + (1 => (node (0 => (node #[13])) (1 => (node #[14])))))))) +-/ +#guard_msgs in +#eval do + let t : DiscrTree Nat := {} + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkApp (mkConst ``h) (mkNatLit 0)) (mkNatLit 0))) 1 + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkApp (mkConst ``h) (mkNatLit 0)) (mkNatLit 1))) 2 + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkApp (mkConst ``h) (mkNatLit 1)) (mkNatLit 0))) 3 + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkApp (mkConst ``h) (mkNatLit 1)) (mkNatLit 1))) 4 + logInfo m!"{t.mapArrays (·.map (· + 10))}" + +/-- +info: ("A" => (node #[10, 11])) +("B" => (node #[12])) +(g => (node (0 => (node #[13, 14, 15])))) +(f => (node (0 => (node #[16])) (1 => (node #[17, 18])) (2 => (node #[19])))) +--- +info: ("A" => (node #[10, 11])) (g => (node (0 => (node #[13, 14, 15])))) (f => (node (1 => (node #[17, 18])))) +--- +info: ("A" => (node #[10])) (g => (node (0 => (node #[14])))) (f => (node (1 => (node #[18])))) +--- +info: +-/ +#guard_msgs in +#eval do + let t : DiscrTree Nat := {} + let t ← t.insert (mkStrLit "A") 0 + let t ← t.insert (mkStrLit "A") 1 + let t ← t.insert (mkStrLit "B") 2 + let t ← t.insert (mkApp (mkConst ``g) (mkNatLit 0)) 3 + let t ← t.insert (mkApp (mkConst ``g) (mkNatLit 0)) 4 + let t ← t.insert (mkApp (mkConst ``g) (mkNatLit 0)) 5 + let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 0)) 6 + let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 1)) 7 + let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 1)) 8 + let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 2)) 9 + logInfo m!"{t.mapArrays (·.map (· + 10))}" + let t := t.mapArrays (fun arr => if arr.size = 1 then #[] else arr) + logInfo m!"{t.mapArrays (·.map (· + 10))}" + let t := t.mapArrays (fun arr => arr.filter (· % 2 = 0)) + logInfo m!"{t.mapArrays (·.map (· + 10))}" + let t := t.mapArrays (fun _ => #[]) + logInfo m!"{t.mapArrays (·.map (· + 10))}" From 6999de194a5193f137443482e2bcacba501616a9 Mon Sep 17 00:00:00 2001 From: Rob Simmons Date: Wed, 19 Aug 2026 13:05:27 -0400 Subject: [PATCH 2/6] review: modify tests and simplify mapArraysM --- src/Lean/Meta/DiscrTree/Util.lean | 6 ++--- tests/elab/discrTreeOps.lean | 44 ++++++++++++++++++++++++------- 2 files changed, 36 insertions(+), 14 deletions(-) diff --git a/src/Lean/Meta/DiscrTree/Util.lean b/src/Lean/Meta/DiscrTree/Util.lean index 1bc8f3c5d8d0..ab98ffb0257c 100644 --- a/src/Lean/Meta/DiscrTree/Util.lean +++ b/src/Lean/Meta/DiscrTree/Util.lean @@ -117,7 +117,7 @@ Checks that a trie node has no values and no children. This is only a check for actual trie emptiness (`t.size = 0`) if all operations maintain the invariant that no trie node has an empty child node. -/ -def Trie.isEmptyNode {α} : Trie α → Bool +def Trie.isEmptyNode : Trie α → Bool | .node vs children => vs.isEmpty && children.isEmpty /-- Apply a monadic function to the array of values at each node in a `DiscrTree`. -/ @@ -137,9 +137,7 @@ partial def Trie.mapArraysM (t : DiscrTree.Trie α) (f : Array α → m (Array /-- Apply a monadic function to the array of values at each node in a `DiscrTree`. -/ def mapArraysM (d : DiscrTree α) (f : Array α → m (Array β)) : m (DiscrTree β) := do let root ← d.root.mapM (fun t => t.mapArraysM f) - let emptyKeys := root.foldl (init := #[]) fun emptyKeys k t => - if t.isEmptyNode then emptyKeys.push k else emptyKeys - pure { root := emptyKeys.foldl (init := root) fun hashMap k => hashMap.erase k } + pure { root := root.foldl (init := root) fun acc k t => if t.isEmptyNode then acc.erase k else acc } /-- Apply a function to the array of values at each node in a `DiscrTree`. -/ def mapArrays (d : DiscrTree α) (f : Array α → Array β) : DiscrTree β := diff --git a/tests/elab/discrTreeOps.lean b/tests/elab/discrTreeOps.lean index 5f1a23fb59ce..5bbc790878b7 100644 --- a/tests/elab/discrTreeOps.lean +++ b/tests/elab/discrTreeOps.lean @@ -79,14 +79,36 @@ 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])))) -(f => (node (0 => (node #[16])) (1 => (node #[17, 18])) (2 => (node #[19])))) +(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])) +("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])))))) --- -info: ("A" => (node #[10, 11])) (g => (node (0 => (node #[13, 14, 15])))) (f => (node (1 => (node #[17, 18])))) +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])))))) --- -info: ("A" => (node #[10])) (g => (node (0 => (node #[14])))) (f => (node (1 => (node #[18])))) +info: +(g => (node (0 => (node #[3, 4, 5])))) +--- +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: -/ @@ -102,11 +124,13 @@ info: let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 0)) 6 let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 1)) 7 let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 1)) 8 + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkConst ``f) (mkNatLit 1))) 10 + let t ← t.insert (mkApp (mkConst ``f) (mkApp (mkConst ``f) (mkNatLit 2))) 11 let t ← t.insert (mkApp (mkConst ``f) (mkNatLit 2)) 9 logInfo m!"{t.mapArrays (·.map (· + 10))}" - let t := t.mapArrays (fun arr => if arr.size = 1 then #[] else arr) - logInfo m!"{t.mapArrays (·.map (· + 10))}" - let t := t.mapArrays (fun arr => arr.filter (· % 2 = 0)) - logInfo m!"{t.mapArrays (·.map (· + 10))}" - let t := t.mapArrays (fun _ => #[]) - logInfo m!"{t.mapArrays (·.map (· + 10))}" + logInfo m!"{t.mapArrays (·.filter (· <= 10))}" + logInfo m!"{t.mapArrays (·.filter (· % 2 = 0))}" + logInfo m!"{t.mapArrays (fun arr => if arr.size > 2 then arr else #[])}" + logInfo m!"{t.mapArrays (fun arr => if arr.size = 1 then arr else #[])}" + logInfo m!"{t.mapArrays (fun arr => if arr.size = 1 then #[] else arr)}" + logInfo m!"{t.mapArrays (β := String) (fun _ => #[])}" From fec7d9f2ac26f5c5884bdc734bf857c24f26a396 Mon Sep 17 00:00:00 2001 From: Rob Simmons Date: Wed, 19 Aug 2026 14:23:34 -0400 Subject: [PATCH 3/6] add @[specalize] and @[inline] annotations --- src/Lean/Meta/DiscrTree/Util.lean | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/Lean/Meta/DiscrTree/Util.lean b/src/Lean/Meta/DiscrTree/Util.lean index ab98ffb0257c..a223f0213111 100644 --- a/src/Lean/Meta/DiscrTree/Util.lean +++ b/src/Lean/Meta/DiscrTree/Util.lean @@ -12,6 +12,7 @@ namespace Trie /-- 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.node vs children => do @@ -29,7 +30,8 @@ def fold (initialKeys : Array Key) (f : σ → Array Key → α → σ) (init : /-- Monadically fold the values stored in a `Trie`. -/ -partial def foldValuesM [Monad m] (f : σ → α → m σ) : (init : σ) → Trie α → m σ + @[specialize] + partial def foldValuesM [Monad m] (f : σ → α → m σ) : (init : σ) → Trie α → m σ | init, node vs children => do let s ← vs.foldlM (init := init) f children.foldlM (init := s) fun s (_, c) => c.foldValuesM (init := s) f @@ -117,10 +119,12 @@ Checks that a trie node has no values and no children. This is only a check for actual trie emptiness (`t.size = 0`) if all operations maintain the invariant that no trie node has an empty child node. -/ +@[inline] def Trie.isEmptyNode : Trie α → Bool | .node vs children => vs.isEmpty && children.isEmpty /-- Apply a monadic function to the array of values at each node in a `DiscrTree`. -/ +@[specialize] partial def Trie.mapArraysM (t : DiscrTree.Trie α) (f : Array α → m (Array β)) : m (DiscrTree.Trie β) := match t with @@ -135,11 +139,13 @@ partial def Trie.mapArraysM (t : DiscrTree.Trie α) (f : Array α → m (Array return .node vs children /-- Apply a monadic function to the array of values at each node in a `DiscrTree`. -/ +@[inline] def mapArraysM (d : DiscrTree α) (f : Array α → m (Array β)) : m (DiscrTree β) := do let root ← d.root.mapM (fun t => t.mapArraysM f) pure { root := root.foldl (init := root) fun acc k t => if t.isEmptyNode then acc.erase k else acc } /-- Apply a function to the array of values at each node in a `DiscrTree`. -/ +@[inline] def mapArrays (d : DiscrTree α) (f : Array α → Array β) : DiscrTree β := Id.run <| d.mapArraysM fun A => pure (f A) From 016ff41f4c79e81eb5af0cc28b4938795abfdfed Mon Sep 17 00:00:00 2001 From: Rob Simmons Date: Wed, 19 Aug 2026 14:26:43 -0400 Subject: [PATCH 4/6] remove stray spaces --- src/Lean/Meta/DiscrTree/Util.lean | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/Lean/Meta/DiscrTree/Util.lean b/src/Lean/Meta/DiscrTree/Util.lean index a223f0213111..adeb45c0bd27 100644 --- a/src/Lean/Meta/DiscrTree/Util.lean +++ b/src/Lean/Meta/DiscrTree/Util.lean @@ -12,7 +12,7 @@ namespace Trie /-- Monadically fold the keys and values stored in a `Trie`. -/ - @[specialize] +@[specialize] partial def foldM [Monad m] (initialKeys : Array Key) (f : σ → Array Key → α → m σ) : (init : σ) → Trie α → m σ | init, Trie.node vs children => do @@ -30,8 +30,8 @@ def fold (initialKeys : Array Key) (f : σ → Array Key → α → σ) (init : /-- Monadically fold the values stored in a `Trie`. -/ - @[specialize] - partial def foldValuesM [Monad m] (f : σ → α → m σ) : (init : σ) → Trie α → m σ +@[specialize] +partial def foldValuesM [Monad m] (f : σ → α → m σ) : (init : σ) → Trie α → m σ | init, node vs children => do let s ← vs.foldlM (init := init) f children.foldlM (init := s) fun s (_, c) => c.foldValuesM (init := s) f From 75d99bade61a5925f6471cd11b4189b85222fe70 Mon Sep 17 00:00:00 2001 From: Rob Simmons Date: Wed, 19 Aug 2026 14:28:18 -0400 Subject: [PATCH 5/6] Improve docstrings --- src/Lean/Meta/DiscrTree/Util.lean | 21 ++++++++++++--------- tests/elab/discrTreeFind.lean | 4 ++++ tests/elab/discrTreeGrind.lean | 2 ++ tests/elab/discrTreeOps.lean | 2 ++ 4 files changed, 20 insertions(+), 9 deletions(-) diff --git a/src/Lean/Meta/DiscrTree/Util.lean b/src/Lean/Meta/DiscrTree/Util.lean index adeb45c0bd27..8b56d1360c40 100644 --- a/src/Lean/Meta/DiscrTree/Util.lean +++ b/src/Lean/Meta/DiscrTree/Util.lean @@ -50,6 +50,16 @@ partial def size : Trie α → Nat | Trie.node vs children => children.foldl (init := vs.size) fun n (_, c) => n + size c +/-- +Checks that a trie node has no values and no children. + +This is only a check for actual trie emptiness (`t.size = 0`) if all operations maintain the +invariant that no trie node has an empty child node. +-/ +@[inline] +def isEmptyNode : Trie α → Bool + | .node vs children => vs.isEmpty && children.isEmpty + end Trie @@ -114,16 +124,9 @@ def size (t : DiscrTree α) : Nat := variable {m : Type → Type} [Monad m] /-- -Checks that a trie node has no values and no children. - -This is only a check for actual trie emptiness (`t.size = 0`) if all operations maintain the -invariant that no trie node has an empty child node. +Apply a monadic function to the array of values at each node in a `DiscrTree`. +Any resulting subtrees containing no values will be pruned. -/ -@[inline] -def Trie.isEmptyNode : Trie α → Bool - | .node vs children => vs.isEmpty && children.isEmpty - -/-- Apply a monadic function to the array of values at each node in a `DiscrTree`. -/ @[specialize] partial def Trie.mapArraysM (t : DiscrTree.Trie α) (f : Array α → m (Array β)) : m (DiscrTree.Trie β) := diff --git a/tests/elab/discrTreeFind.lean b/tests/elab/discrTreeFind.lean index b79340251b2d..7fbe175fbec8 100644 --- a/tests/elab/discrTreeFind.lean +++ b/tests/elab/discrTreeFind.lean @@ -1,3 +1,7 @@ +/-! +Test basic lookup operations (match, match-liberal, and unify) on discrimination trees. +-/ + import Lean open Lean Meta diff --git a/tests/elab/discrTreeGrind.lean b/tests/elab/discrTreeGrind.lean index fa1abebbf241..6b6e198d08be 100644 --- a/tests/elab/discrTreeGrind.lean +++ b/tests/elab/discrTreeGrind.lean @@ -1,3 +1,5 @@ +/-! Exercise grind's usage of DiscrTrees -/ + def F (x : Nat) : Nat := x def G (x : Nat) : Nat := x def H (x : Nat) : Nat := x diff --git a/tests/elab/discrTreeOps.lean b/tests/elab/discrTreeOps.lean index 5bbc790878b7..e73d6d0baaab 100644 --- a/tests/elab/discrTreeOps.lean +++ b/tests/elab/discrTreeOps.lean @@ -1,3 +1,5 @@ +/-! Exercise basic operations on discrimination trees -/ + import Lean open Lean Meta From 1be6c1fb31abdc19c1c9a56276dc9e3f9cb18fe1 Mon Sep 17 00:00:00 2001 From: Rob Simmons Date: Wed, 19 Aug 2026 15:44:00 -0400 Subject: [PATCH 6/6] fix test case docstrings --- tests/elab/discrTreeFind.lean | 3 ++- tests/elab/discrTreeOps.lean | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/elab/discrTreeFind.lean b/tests/elab/discrTreeFind.lean index 7fbe175fbec8..b33afa78449b 100644 --- a/tests/elab/discrTreeFind.lean +++ b/tests/elab/discrTreeFind.lean @@ -1,8 +1,9 @@ +import Lean + /-! Test basic lookup operations (match, match-liberal, and unify) on discrimination trees. -/ -import Lean open Lean Meta opaque a : Nat diff --git a/tests/elab/discrTreeOps.lean b/tests/elab/discrTreeOps.lean index e73d6d0baaab..5165da64a8f1 100644 --- a/tests/elab/discrTreeOps.lean +++ b/tests/elab/discrTreeOps.lean @@ -1,6 +1,7 @@ +import Lean + /-! Exercise basic operations on discrimination trees -/ -import Lean open Lean Meta opaque f : Nat → Nat