feat: add insertPattern for discrimination tree insertion in Sym (#11884)

This PR adds discrimination tree support for the symbolic simulation
framework.
The new `DiscrTree.lean` module converts `Pattern` values into
discrimination
tree keys, treating proof/instance arguments and pattern variables as
wildcards
(`Key.star`). Motivation: efficient pattern retrieval during rewriting.

src/Lean/Elab/Tactic/Do/Attr.lean

@@ -193,7 +193,7 @@ def mkSpecTheoremFromStx (ref : Syntax) (proof : Expr) (prio : Nat := eval_prio
   mkSpecTheorem type (.stx (← mkFreshId) ref proof) prio
 
 def addSpecTheoremEntry (d : SpecTheorems) (e : SpecTheorem) : SpecTheorems :=
-  { d with specs := d.specs.insertCore e.keys e }
+  { d with specs := d.specs.insertKeyValue e.keys e }
 
 def addSpecTheorem (ext : SpecExtension) (declName : Name) (prio : Nat) (attrKind : AttributeKind) : MetaM Unit := do
   let thm ← mkSpecTheoremFromConst declName prio

src/Lean/Elab/Tactic/Monotonicity.lean

@@ -29,7 +29,7 @@ partial def headBetaUnderLambda (f : Expr) : Expr := Id.run do
 builtin_initialize monotoneExt :
     SimpleScopedEnvExtension (Name × Array DiscrTree.Key) (DiscrTree Name) ←
   registerSimpleScopedEnvExtension {
-    addEntry := fun dt (n, ks) => dt.insertCore ks n
+    addEntry := fun dt (n, ks) => dt.insertKeyValue ks n
     initial := {}
   }
 

src/Lean/Meta/DiscrTree/Basic.lean

@@ -11,6 +11,12 @@ public import Lean.CoreM
 public section
 namespace Lean.Meta.DiscrTree
 
+def mkNoindexAnnotation (e : Expr) : Expr :=
+  mkAnnotation `noindex e
+
+def hasNoindexAnnotation (e : Expr) : Bool :=
+  annotation? `noindex e |>.isSome
+
 instance : Inhabited (Trie α) where
   default := .node #[] #[]
 
@@ -115,4 +121,60 @@ where
       r := r.push (← go)
     return r
 
+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)]
+  else
+    .node #[v] #[]
+
+/--
+If `vs` contains an element `v'` such that `v == v'`, then replace `v'` with `v`.
+Otherwise, push `v`.
+See issue #2155
+Recall that `BEq α` may not be Lawful.
+-/
+private def insertVal [BEq α] (vs : Array α) (v : α) : Array α :=
+  loop 0
+where
+  loop (i : Nat) : Array α :=
+    if h : i < vs.size then
+      if v == vs[i] then
+        vs.set i v
+      else
+        loop (i+1)
+    else
+      vs.push v
+  termination_by vs.size - i
+
+private partial def insertAux [BEq α] (keys : Array Key) (v : α) : Nat → Trie α → Trie α
+  | i, .node vs cs =>
+    if h : i < keys.size then
+      let k := keys[i]
+      let c := Id.run $ cs.binInsertM
+          (fun a b => a.1 < b.1)
+          (fun ⟨_, s⟩ => let c := insertAux keys v (i+1) s; (k, c)) -- merge with existing
+          (fun _ => let c := createNodes keys v (i+1); (k, c))
+          (k, default)
+      .node vs c
+    else
+      .node (insertVal vs v) cs
+
+def insertKeyValue [BEq α] (d : DiscrTree α) (keys : Array Key) (v : α) : DiscrTree α :=
+  if keys.isEmpty then panic! "invalid key sequence"
+  else
+    let k := keys[0]!
+    match d.root.find? k with
+    | none =>
+      let c := createNodes keys v 1
+      { root := d.root.insert k c }
+    | some c =>
+      let c := insertAux keys v 1 c
+      { root := d.root.insert k c }
+
+@[deprecated insertKeyValue (since := "2026-01-02")]
+def insertCore [BEq α] (d : DiscrTree α) (keys : Array Key) (v : α) : DiscrTree α :=
+  insertKeyValue d keys v
+
 end Lean.Meta.DiscrTree

src/Lean/Meta/DiscrTree/Main.lean

@@ -198,12 +198,6 @@ private def isOffset (fName : Name) (e : Expr) : MetaM Bool := do
 private def shouldAddAsStar (fName : Name) (e : Expr) : MetaM Bool := do
   isOffset fName e
 
-def mkNoindexAnnotation (e : Expr) : Expr :=
-  mkAnnotation `noindex e
-
-def hasNoindexAnnotation (e : Expr) : Bool :=
-  annotation? `noindex e |>.isSome
-
 /--
 Reduction procedure for the discrimination tree indexing.
 -/
@@ -338,61 +332,9 @@ def mkPath (e : Expr) (noIndexAtArgs := false) : MetaM (Array Key) := do
     let keys : Array Key := .mkEmpty initCapacity
     mkPathAux (root := true) (todo.push e) keys noIndexAtArgs
 
-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)]
-  else
-    .node #[v] #[]
-
-/--
-If `vs` contains an element `v'` such that `v == v'`, then replace `v'` with `v`.
-Otherwise, push `v`.
-See issue #2155
-Recall that `BEq α` may not be Lawful.
--/
-private def insertVal [BEq α] (vs : Array α) (v : α) : Array α :=
-  loop 0
-where
-  loop (i : Nat) : Array α :=
-    if h : i < vs.size then
-      if v == vs[i] then
-        vs.set i v
-      else
-        loop (i+1)
-    else
-      vs.push v
-  termination_by vs.size - i
-
-private partial def insertAux [BEq α] (keys : Array Key) (v : α) : Nat → Trie α → Trie α
-  | i, .node vs cs =>
-    if h : i < keys.size then
-      let k := keys[i]
-      let c := Id.run $ cs.binInsertM
-          (fun a b => a.1 < b.1)
-          (fun ⟨_, s⟩ => let c := insertAux keys v (i+1) s; (k, c)) -- merge with existing
-          (fun _ => let c := createNodes keys v (i+1); (k, c))
-          (k, default)
-      .node vs c
-    else
-      .node (insertVal vs v) cs
-
-def insertCore [BEq α] (d : DiscrTree α) (keys : Array Key) (v : α) : DiscrTree α :=
-  if keys.isEmpty then panic! "invalid key sequence"
-  else
-    let k := keys[0]!
-    match d.root.find? k with
-    | none =>
-      let c := createNodes keys v 1
-      { root := d.root.insert k c }
-    | some c =>
-      let c := insertAux keys v 1 c
-      { root := d.root.insert k c }
-
 def insert [BEq α] (d : DiscrTree α) (e : Expr) (v : α) (noIndexAtArgs := false) : MetaM (DiscrTree α) := do
   let keys ← mkPath e noIndexAtArgs
-  return d.insertCore keys v
+  return d.insertKeyValue keys v
 
 /--
 Inserts a value into a discrimination tree,
@@ -400,10 +342,10 @@ but only if its key is not of the form `#[*]` or `#[=, *, *, *]`.
 -/
 def insertIfSpecific [BEq α] (d : DiscrTree α) (e : Expr) (v : α) (noIndexAtArgs := false) : MetaM (DiscrTree α) := do
   let keys ← mkPath e noIndexAtArgs
-  return if keys == #[Key.star] || keys == #[Key.const `Eq 3, Key.star, Key.star, Key.star] then
-    d
+  if keys == #[Key.star] || keys == #[Key.const `Eq 3, Key.star, Key.star, Key.star] then
+    return d
   else
-    d.insertCore keys v
+    return d.insertKeyValue keys v
 
 private def getKeyArgs (e : Expr) (isMatch root : Bool) : MetaM (Key × Array Expr) := do
   let e ← reduceDT e root

src/Lean/Meta/DiscrTree/Types.lean

@@ -5,7 +5,7 @@ Authors: Leonardo de Moura
 -/
 module
 prelude
-public import Lean.ToExpr
+public import Lean.Expr
 public section
 namespace Lean.Meta
 /-! See file `DiscrTree.lean` for the actual implementation and documentation. -/

src/Lean/Meta/Instances.lean

@@ -75,8 +75,8 @@ structure Instances where
 
 def addInstanceEntry (d : Instances) (e : InstanceEntry) : Instances :=
   match e.globalName? with
-  | some n => { d with discrTree := d.discrTree.insertCore e.keys e, instanceNames := d.instanceNames.insert n e, erased := d.erased.erase n }
-  | none   => { d with discrTree := d.discrTree.insertCore e.keys e }
+  | some n => { d with discrTree := d.discrTree.insertKeyValue e.keys e, instanceNames := d.instanceNames.insert n e, erased := d.erased.erase n }
+  | none   => { d with discrTree := d.discrTree.insertKeyValue e.keys e }
 
 def Instances.eraseCore (d : Instances) (declName : Name) : Instances :=
   { d with erased := d.erased.insert declName, instanceNames := d.instanceNames.erase declName }

src/Lean/Meta/Sym.lean

@@ -26,6 +26,7 @@ public import Lean.Meta.Sym.EqTrans
 public import Lean.Meta.Sym.Congr
 public import Lean.Meta.Sym.SimpResult
 public import Lean.Meta.Sym.Simp
+public import Lean.Meta.Sym.DiscrTree
 
 /-!
 # Symbolic simulation support.

src/Lean/Meta/Sym/DiscrTree.lean

@@ -0,0 +1,139 @@
+/-
+Copyright (c) 2026 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+module
+prelude
+public import Lean.Meta.Sym.Pattern
+import Lean.Meta.DiscrTree.Basic
+namespace Lean.Meta.Sym
+open DiscrTree
+
+/-!
+# Discrimination Tree for the Symbolic Simulation Framework
+
+This module provides pattern insertion into discrimination trees for the symbolic simulation
+framework. Unlike the standard `DiscrTree` insertion which works with arbitrary expressions,
+this module converts `Pattern` values (see `Sym/Pattern.lean`) into discrimination tree keys.
+`Pattern` values have been preprocessed and use de Bruijn variables for wildcards instead of
+metavariables.
+
+## Key Design Decisions
+
+The conversion from `Pattern` to `Array Key` respects the structural matching semantics of `Sym`:
+
+1. **Proof and instance arguments are wildcards.** When building keys for a function application,
+   arguments marked as proofs or instances in `ProofInstInfo` are replaced with `Key.star`.
+   This ensures that during retrieval, terms differing only in proof/instance arguments will
+   match the same entries.
+
+2. **Bound variables are wildcards.** Pattern variables (represented as de Bruijn indices in
+   `Pattern.pattern`) become `Key.star`, allowing them to match any subterm during retrieval.
+
+3. **`noindex` annotations are respected.** Subterms annotated with `noindex` become `Key.star`.
+
+## Usage
+
+Use `insertPattern` to add a pattern to a discrimination tree:
+```
+let tree := insertPattern tree pattern value
+```
+
+Retrieval should use the standard `DiscrTree.getMatch` or similar, which will find all patterns
+whose key sequence is compatible with the query term.
+-/
+
+/--
+Returns the number of child keys for a given discrimination tree key.
+**Note**: Unlike the standard `DiscrTree` module, `Key.arrow` has arity 2.
+-/
+def getKeyArity : Key → Nat
+  | .const _ a  => a
+  | .fvar _ a   => a
+  | .arrow      => 2
+  | _           => 0
+
+/-- Returns `true` if argument at position `i` should be ignored (is a proof or instance). -/
+def ignoreArg (infos : Array ProofInstArgInfo) (i : Nat) : Bool :=
+  if h : i < infos.size then
+    let info := infos[i]
+    info.isInstance || info.isProof
+  else
+    false
+
+/-- Pushes all arguments of an application onto the todo stack in reverse order. -/
+def pushAllArgs (e : Expr) (todo : Array Expr) : Array Expr :=
+  match e with
+  | .app f a => pushAllArgs f (todo.push a)
+  | _ => todo
+
+/--
+Pushes arguments of an application onto the todo stack, replacing proof/instance arguments
+with dummy bound variables (which become `Key.star` wildcards).
+-/
+def pushArgsUsingInfo (infos : Array ProofInstArgInfo) (i : Nat) (e : Expr) (todo : Array Expr) : Array Expr :=
+  match e with
+  | .app f a =>
+    if ignoreArg infos i then
+      -- **Note**: We use a dummy bvar, it will be transformed into a `Key.star` later.
+      let dummyBVar := .bvar 1000000
+      pushArgsUsingInfo infos (i-1) f (todo.push dummyBVar)
+    else
+      pushArgsUsingInfo infos (i-1) f (todo.push a)
+  | _ => todo
+
+/--
+Computes the discrimination tree key for an expression and pushes its subterms onto the todo stack.
+Returns `Key.star` for bound variables and `noindex`-annotated terms.
+-/
+def pushArgs (fnInfos : AssocList Name ProofInstInfo) (todo : Array Expr) (e : Expr) : Key × Array Expr :=
+  if hasNoindexAnnotation e then
+    (.star, todo)
+  else
+    let fn := e.getAppFn
+    match fn with
+    | .lit v => (.lit v, todo)
+    | .bvar _ => (.star, todo)
+    | .forallE _ d b _ => (.arrow, todo.push b |>.push d)
+    | .const declName _ =>
+      let numArgs := e.getAppNumArgs
+      let todo := if let some info := fnInfos.find? declName then
+        pushArgsUsingInfo info.argsInfo (numArgs - 1) e todo
+      else
+        pushAllArgs e todo
+      (.const declName numArgs, todo)
+    | .fvar fvarId   =>
+      let numArgs := e.getAppNumArgs
+      let todo := pushAllArgs e todo
+      (.fvar fvarId numArgs, todo)
+    | _ => (.other, todo)
+
+/-- Work-list based traversal that builds the key sequence for a pattern. -/
+partial def mkPathAux (fnInfos : AssocList Name ProofInstInfo) (todo : Array Expr) (keys : Array Key) : Array Key :=
+  if todo.isEmpty then
+    keys
+  else
+    let e    := todo.back!
+    let todo := todo.pop
+    let (k, todo) := pushArgs fnInfos todo e
+    mkPathAux fnInfos todo (keys.push k)
+
+def initCapacity := 8
+
+/-- Converts a `Pattern` into a discrimination tree key sequence. -/
+public def Pattern.mkDiscrTreeKeys (p : Pattern) : Array Key :=
+  let todo : Array Expr := .mkEmpty initCapacity
+  let keys : Array Key := .mkEmpty initCapacity
+  mkPathAux p.fnInfos (todo.push p.pattern) keys
+
+/-- Inserts a pattern into a discrimination tree, associating it with value `v`. -/
+public def insertPattern [BEq α] (d : DiscrTree α) (p : Pattern) (v : α) : DiscrTree α :=
+  let keys := p.mkDiscrTreeKeys
+  d.insertKeyValue keys v
+
+/-!
+**TODO** Retrieval.
+-/
+
+end Lean.Meta.Sym

src/Lean/Meta/Tactic/Ext.lean

@@ -40,7 +40,7 @@ builtin_initialize extExtension :
     SimpleScopedEnvExtension ExtTheorem ExtTheorems ←
   registerSimpleScopedEnvExtension {
     addEntry := fun { tree, erased } thm =>
-      { tree := tree.insertCore thm.keys thm, erased := erased.erase thm.declName }
+      { tree := tree.insertKeyValue thm.keys thm, erased := erased.erase thm.declName }
     initial := {}
   }
 

src/Lean/Meta/Tactic/Rfl.lean

@@ -26,7 +26,7 @@ open Lean Meta
 initialize reflExt :
     SimpleScopedEnvExtension (Name × Array DiscrTree.Key) (DiscrTree Name) ←
   registerSimpleScopedEnvExtension {
-    addEntry := fun dt (n, ks) => dt.insertCore ks n
+    addEntry := fun dt (n, ks) => dt.insertKeyValue ks n
     initial := {}
   }
 

src/Lean/Meta/Tactic/Simp/SimpTheorems.lean

@@ -572,9 +572,9 @@ def SimpTheorems.addSimpTheorem (d : SimpTheorems) (e : SimpTheorem) : SimpTheor
   -- Erase the converse, if it exists
   let d := eraseFwdIfBwd d e
   if e.post then
-    { d with post := d.post.insertCore e.keys e, lemmaNames := updateLemmaNames d.lemmaNames }
+    { d with post := d.post.insertKeyValue e.keys e, lemmaNames := updateLemmaNames d.lemmaNames }
   else
-    { d with pre := d.pre.insertCore e.keys e, lemmaNames := updateLemmaNames d.lemmaNames }
+    { d with pre := d.pre.insertKeyValue e.keys e, lemmaNames := updateLemmaNames d.lemmaNames }
 where
   updateLemmaNames (s : PHashSet Origin) : PHashSet Origin :=
     s.insert e.origin

src/Lean/Meta/Tactic/Simp/Simproc.lean

@@ -153,9 +153,9 @@ def addSimprocAttrCore (ext : SimprocExtension) (declName : Name) (kind : Attrib
 def Simprocs.addCore (s : Simprocs) (keys : Array SimpTheoremKey) (declName : Name) (post : Bool) (proc : Sum Simproc DSimproc) : Simprocs :=
   let s := { s with simprocNames := s.simprocNames.insert declName, erased := s.erased.erase declName }
   if post then
-    { s with post := s.post.insertCore keys { declName, keys, post, proc } }
+    { s with post := s.post.insertKeyValue keys { declName, keys, post, proc } }
   else
-    { s with pre := s.pre.insertCore keys { declName, keys, post, proc } }
+    { s with pre := s.pre.insertKeyValue keys { declName, keys, post, proc } }
 
 /--
 Implements attributes `builtin_simproc` and `builtin_sevalproc`.

src/Lean/Meta/Tactic/Symm.lean

@@ -26,7 +26,7 @@ namespace Lean.Meta.Symm
 builtin_initialize symmExt :
     SimpleScopedEnvExtension (Name × Array DiscrTree.Key) (DiscrTree Name) ←
   registerSimpleScopedEnvExtension {
-    addEntry := fun dt (n, ks) => dt.insertCore ks n
+    addEntry := fun dt (n, ks) => dt.insertKeyValue ks n
     initial := {}
   }
 

src/Lean/Meta/UnificationHint.lean

@@ -32,7 +32,7 @@ private def config : ConfigWithKey :=
   { iota := false, proj := .no : Config }.toConfigWithKey
 
 def UnificationHints.add (hints : UnificationHints) (e : UnificationHintEntry) : UnificationHints :=
-  { hints with discrTree := hints.discrTree.insertCore e.keys e.val }
+  { hints with discrTree := hints.discrTree.insertKeyValue e.keys e.val }
 
 builtin_initialize unificationHintExtension : SimpleScopedEnvExtension UnificationHintEntry UnificationHints ←
   registerSimpleScopedEnvExtension {

tests/lean/run/sym_pattern_2.lean

@@ -1,4 +1,5 @@
 import Lean.Meta.Sym
+import Lean.Meta.DiscrTree.Basic
 open Lean Meta Sym Grind
 set_option grind.debug true
 opaque p [Ring α] : α → α → Prop
@@ -58,3 +59,65 @@ info: ∀ (x : Nat), q (f (f x)) (f x + f (f 1))
 -/
 #guard_msgs in
 #eval SymM.run' test₂
+
+theorem forall_and_eq (P Q : α → Prop) : (∀ x, P x ∧ Q x) = ((∀ x, P x) ∧ (∀ x, Q x)):= by
+  grind
+
+def logPatternKey (p : Pattern) : MetaM Unit := do
+  let k := p.mkDiscrTreeKeys
+  logInfo m!"{k.toList.map (·.format)}"
+
+def logPatternKeyFor (declName : Name) : MetaM Unit := do
+  let (p, _) ← mkEqPatternFromDecl declName
+  logPatternKey p
+
+/--
+info: [HAdd.hAdd, Nat, Nat, Nat, *, OfNat.ofNat, Nat, 0, *, *]
+---
+info: [HMul.hMul, *, *, *, *, OfNat.ofNat, *, 0, *, *]
+---
+info: [∀, *, And, *, *]
+---
+info: [Array.eraseIdx, *, HAppend.hAppend, Array, *, Array, *, Array, *, *, *, *, *, *]
+---
+info: [List.map, *, *, *, List.map, *, *, *, *]
+---
+info: [Std.HashMap.insertMany,
+ *,
+ *,
+ *,
+ *,
+ List,
+ Prod,
+ *,
+ *,
+ *,
+ *,
+ HAppend.hAppend,
+ List,
+ Prod,
+ *,
+ *,
+ List,
+ Prod,
+ *,
+ *,
+ List,
+ Prod,
+ *,
+ *,
+ *,
+ *,
+ *]
+---
+info: [GetElem.getElem, Std.HashMap, *, *, *, *, *, *, ◾, *, Std.HashMap.insert, *, *, *, *, *, *, *, *, *]
+-/
+#guard_msgs in
+#eval SymM.run' do
+  logPatternKeyFor ``Nat.zero_add
+  logPatternKeyFor ``Grind.Semiring.zero_mul
+  logPatternKeyFor ``forall_and_eq
+  logPatternKeyFor ``Array.eraseIdx_append
+  logPatternKeyFor ``List.map_map
+  logPatternKeyFor ``Std.HashMap.insertMany_append
+  logPatternKeyFor ``Std.HashMap.getElem_insert