diff --git a/src/Lean/Meta/DiscrTree/Util.lean b/src/Lean/Meta/DiscrTree/Util.lean index 1916e78a59cd..8b56d1360c40 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,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) diff --git a/tests/elab/discrTreeFind.lean b/tests/elab/discrTreeFind.lean new file mode 100644 index 000000000000..b33afa78449b --- /dev/null +++ b/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))))}" diff --git a/tests/elab/discrTreeGrind.lean b/tests/elab/discrTreeGrind.lean new file mode 100644 index 000000000000..6b6e198d08be --- /dev/null +++ b/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 diff --git a/tests/elab/discrTreeOps.lean b/tests/elab/discrTreeOps.lean new file mode 100644 index 000000000000..5165da64a8f1 --- /dev/null +++ b/tests/elab/discrTreeOps.lean @@ -0,0 +1,139 @@ +import Lean + +/-! Exercise basic operations on discrimination trees -/ + +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])) + (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 #[0])) +("B" => (node #[2])) +(g => (node (0 => (node #[4])))) +(f => (node (0 => (node #[6])) (1 => (node #[8])) (f => (node (1 => (node #[10])))))) +--- +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: +-/ +#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) (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))}" + 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 _ => #[])}"