fix: collapse empty DiscrTree tries

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,6 +30,7 @@ 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 σ
   | init, node vs children => do
     let s ← vs.foldlM (init := init) f
@@ -48,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
 
 
@@ -111,18 +123,32 @@ def size (t : DiscrTree α) : Nat :=
 
 variable {m : Type → Type} [Monad m]
 
-/-- Apply a monadic function to the array of values at each node in a `DiscrTree`. -/
+/--
+Apply a monadic function to the array of values at each node in a `DiscrTree`.
+Any resulting subtrees containing no values will be pruned.
+-/
+@[specialize]
 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`. -/
+@[inline]
 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)
+  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)
 

tests/elab/discrTreeFind.lean

@@ -0,0 +1,210 @@
+import Lean
+
+/-!
+Test basic lookup operations (match, match-liberal, and unify) on discrimination trees.
+-/
+
+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))))}"

tests/elab/discrTreeGrind.lean

@@ -0,0 +1,11 @@
+/-! Exercise grind's usage of DiscrTrees -/
+
+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