From cce5c33a4a0ef2bd199b4dafb85fcbf4bcbabf5b Mon Sep 17 00:00:00 2001
From: Rob Simmons <rob@lean-fro.org>
Date: Mon, 17 Aug 2026 12:10:31 -0400
Subject: [PATCH 1/8] Add .chain branch without use, panic everywhere

---
 src/Lean/Meta/DiscrTree/Basic.lean    | 4 ++++
 src/Lean/Meta/DiscrTree/Main.lean     | 5 +++++
 src/Lean/Meta/DiscrTree/Types.lean    | 5 +++++
 src/Lean/Meta/DiscrTree/Util.lean     | 4 ++++
 src/Lean/Meta/Sym/Simp/DiscrTree.lean | 2 ++
 5 files changed, 20 insertions(+)

diff --git a/src/Lean/Meta/DiscrTree/Basic.lean b/src/Lean/Meta/DiscrTree/Basic.lean
index 745033421224..bc0e65d59fa6 100644
--- a/src/Lean/Meta/DiscrTree/Basic.lean
+++ b/src/Lean/Meta/DiscrTree/Basic.lean
@@ -60,6 +60,9 @@ def Key.format : Key → Format
 instance : ToFormat Key := ⟨Key.format⟩
 
 partial def Trie.format [ToFormat α] : Trie α → Format
+  | .chain k c =>
+    -- TODO: Make a separate formatting for chains instead of replicating the node formatting
+    Trie.format (.node #[] #[(k, 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))
@@ -150,6 +153,7 @@ where
   termination_by vs.size - i
 
 private partial def insertAux [BEq α] (keys : Array Key) (v : α) : Nat → Trie α → Trie α
+  | i, .chain _k _c => panic! "unimpl"
   | 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..0326da8ceb4e 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 _k _c) => panic! "unimpl"
   | some (.node vs _) => result ++ vs
 
 private abbrev findKey (cs : Array (Key × Trie α)) (k : Key) : Option (Key × Trie α) :=
@@ -437,6 +438,7 @@ 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 _k _c => panic! "unimpl"
   | .node vs cs =>
     if todo.isEmpty then
       return result ++ vs
@@ -544,6 +546,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 _k _c => panic! "unimpl"
     | .node vs cs =>
       let mut result := result ++ vs
       for (_, trie) in cs do
@@ -577,11 +580,13 @@ 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 _k _c => panic! "unimpl"
     | 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 _k _c => panic! "unimpl"
     | 0, .node vs cs => do
       if todo.isEmpty then
         return result ++ vs
diff --git a/src/Lean/Meta/DiscrTree/Types.lean b/src/Lean/Meta/DiscrTree/Types.lean
index 45e29c6b0b24..eb0b037177bb 100644
--- a/src/Lean/Meta/DiscrTree/Types.lean
+++ b/src/Lean/Meta/DiscrTree/Types.lean
@@ -38,6 +38,11 @@ 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. Special-casing 
+    -/
+  | 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..0d572dcdb554 100644
--- a/src/Lean/Meta/DiscrTree/Util.lean
+++ b/src/Lean/Meta/DiscrTree/Util.lean
@@ -15,6 +15,7 @@ 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 => pure init -- UNSOUND
   | 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 +33,7 @@ Monadically fold the values stored in a `Trie`.
 -/
 @[specialize]
 partial def foldValuesM [Monad m] (f : σ → α → m σ) : (init : σ) → Trie α → m σ
+  | init, chain _k _c => pure init -- UNSOUND
   | 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 +49,7 @@ def foldValues (f : σ → α → σ) (init : σ) (t : Trie α) : σ :=
 The number of values stored in a `Trie`.
 -/
 partial def size : Trie α → Nat
+  | Trie.chain _k _c => panic! "unimpl"
   | Trie.node vs children =>
     children.foldl (init := vs.size) fun n (_, c) => n + size c
 
@@ -161,6 +164,7 @@ 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 => panic! "unimpl"
   | .node vs children => do
     let vs ← f vs
     let children ← children.filterMapM fun (k, child) => do
diff --git a/src/Lean/Meta/Sym/Simp/DiscrTree.lean b/src/Lean/Meta/Sym/Simp/DiscrTree.lean
index 2152e1a9a0c9..7c010ac9e64a 100644
--- a/src/Lean/Meta/Sym/Simp/DiscrTree.lean
+++ b/src/Lean/Meta/Sym/Simp/DiscrTree.lean
@@ -167,6 +167,7 @@ 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 k c => panic! "unimpl"
   | .node vs cs =>
     let csize := cs.size
     if todo.isEmpty then
@@ -204,6 +205,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 _ _) => panic! "unimpl"
   | some (.node vs _) => vs
   let e := resolveAssignedMVars mctx <| etaReduce e
   match d.root.find? (getKey e) with

From 71322d467cece5c9acc6bfa25f3a046a13c14b08 Mon Sep 17 00:00:00 2001
From: Rob Simmons <robsimmons@gmail.com>
Date: Mon, 17 Aug 2026 12:17:21 -0400
Subject: [PATCH 2/8] introduce chain nodes; minimal changes to make some tests
 work

---
 src/Lean/Meta/DiscrTree/Basic.lean    | 13 +++++++++++--
 src/Lean/Meta/DiscrTree/Main.lean     |  9 ++++++---
 src/Lean/Meta/DiscrTree/Util.lean     |  9 ++++++++-
 src/Lean/Meta/Sym/Simp/DiscrTree.lean |  5 ++++-
 4 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/src/Lean/Meta/DiscrTree/Basic.lean b/src/Lean/Meta/DiscrTree/Basic.lean
index bc0e65d59fa6..54aa59b10940 100644
--- a/src/Lean/Meta/DiscrTree/Basic.lean
+++ b/src/Lean/Meta/DiscrTree/Basic.lean
@@ -129,7 +129,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] #[]
 
@@ -153,7 +153,16 @@ where
   termination_by vs.size - i
 
 private partial def insertAux [BEq α] (keys : Array Key) (v : α) : Nat → Trie α → Trie α
-  | i, .chain _k _c => panic! "unimpl"
+  | 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 0326da8ceb4e..3954c4300e62 100644
--- a/src/Lean/Meta/DiscrTree/Main.lean
+++ b/src/Lean/Meta/DiscrTree/Main.lean
@@ -438,7 +438,9 @@ 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 _k _c => panic! "unimpl"
+  | .chain k c =>
+    -- Reuse general code path
+    getMatchLoop todo (.node #[] #[(k, c)]) result
   | .node vs cs =>
     if todo.isEmpty then
       return result ++ vs
@@ -580,13 +582,14 @@ 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 _k _c => panic! "unimpl"
+    | _, .chain k c =>
+      -- Reuse general code path
+      process skip todo (.node #[] #[(k, c)]) 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 _k _c => panic! "unimpl"
     | 0, .node vs cs => do
       if todo.isEmpty then
         return result ++ vs
diff --git a/src/Lean/Meta/DiscrTree/Util.lean b/src/Lean/Meta/DiscrTree/Util.lean
index 0d572dcdb554..7641aa6ef64b 100644
--- a/src/Lean/Meta/DiscrTree/Util.lean
+++ b/src/Lean/Meta/DiscrTree/Util.lean
@@ -58,13 +58,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⟩
 
 /--
@@ -73,6 +77,7 @@ Equivalent to `t.asNode.1`.
 -/
 @[inline]
 def nodeValues : Trie α → Array α
+  | .chain _ _ => #[]
   | .node vs _ => vs
 
 /--
@@ -81,6 +86,7 @@ Equivalent to `t.asNode.2`.
 -/
 @[inline]
 def nodeChildren : Trie α → Array (Key × Trie α)
+  | .chain k v => #[(k, v)]
   | .node _ cs => cs
 
 /--
@@ -91,6 +97,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
diff --git a/src/Lean/Meta/Sym/Simp/DiscrTree.lean b/src/Lean/Meta/Sym/Simp/DiscrTree.lean
index 7c010ac9e64a..78e6569d9133 100644
--- a/src/Lean/Meta/Sym/Simp/DiscrTree.lean
+++ b/src/Lean/Meta/Sym/Simp/DiscrTree.lean
@@ -167,7 +167,10 @@ 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 k c => panic! "unimpl"
+  | .chain k c =>
+    -- Reuse general code path
+    getMatchLoop mctx todo (.node #[] #[(k, c)]) result
+
   | .node vs cs =>
     let csize := cs.size
     if todo.isEmpty then

From aad1a068cc6665b0a2b973cbd1f042a4fdc1ca5c Mon Sep 17 00:00:00 2001
From: Rob Simmons <robsimmons@gmail.com>
Date: Mon, 17 Aug 2026 14:57:55 -0400
Subject: [PATCH 3/8] fix bugs demanded by test cases

more tests, mapArrays fix and test cases


better test coverage for basic discrimination tree options


refactor getMatchLoop(s) for .chain


Add public interface Trie.mkNode / Trie.asNode for downstream's sake


formatting
---
 src/Lean/Meta/DiscrTree/Main.lean     | 33 ++++++++++++++++--------
 src/Lean/Meta/DiscrTree/Util.lean     | 10 +++++---
 src/Lean/Meta/Sym/Simp/DiscrTree.lean | 37 ++++++++++++++-------------
 3 files changed, 47 insertions(+), 33 deletions(-)

diff --git a/src/Lean/Meta/DiscrTree/Main.lean b/src/Lean/Meta/DiscrTree/Main.lean
index 3954c4300e62..8803a288ed18 100644
--- a/src/Lean/Meta/DiscrTree/Main.lean
+++ b/src/Lean/Meta/DiscrTree/Main.lean
@@ -438,9 +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 k c =>
-    -- Reuse general code path
-    getMatchLoop todo (.node #[] #[(k, c)]) result
+  | .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
@@ -454,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
@@ -548,7 +558,8 @@ private partial def getAllValuesForKey (d : DiscrTree α) (k : Key) (result : Ar
 where
   go (trie : Trie α) (result : Array α) : Array α := Id.run do
     match trie with
-    | .chain _k _c => panic! "unimpl"
+    | .chain k c =>
+      go (.node #[] #[(k, c)]) result
     | .node vs cs =>
       let mut result := result ++ vs
       for (_, trie) in cs do
diff --git a/src/Lean/Meta/DiscrTree/Util.lean b/src/Lean/Meta/DiscrTree/Util.lean
index 7641aa6ef64b..44c64e9cfac0 100644
--- a/src/Lean/Meta/DiscrTree/Util.lean
+++ b/src/Lean/Meta/DiscrTree/Util.lean
@@ -15,7 +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 => pure init -- UNSOUND
+  | 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) =>
@@ -33,7 +34,8 @@ Monadically fold the values stored in a `Trie`.
 -/
 @[specialize]
 partial def foldValuesM [Monad m] (f : σ → α → m σ) : (init : σ) → Trie α → m σ
-  | init, chain _k _c => pure init -- UNSOUND
+  | 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
@@ -49,7 +51,7 @@ def foldValues (f : σ → α → σ) (init : σ) (t : Trie α) : σ :=
 The number of values stored in a `Trie`.
 -/
 partial def size : Trie α → Nat
-  | Trie.chain _k _c => panic! "unimpl"
+  | Trie.chain _ c => size c
   | Trie.node vs children =>
     children.foldl (init := vs.size) fun n (_, c) => n + size c
 
@@ -171,7 +173,7 @@ 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 => panic! "unimpl"
+  | .chain k c => Trie.mapArraysM (.node #[] #[(k, c)]) f
   | .node vs children => do
     let vs ← f vs
     let children ← children.filterMapM fun (k, child) => do
diff --git a/src/Lean/Meta/Sym/Simp/DiscrTree.lean b/src/Lean/Meta/Sym/Simp/DiscrTree.lean
index 78e6569d9133..135be8aa1817 100644
--- a/src/Lean/Meta/Sym/Simp/DiscrTree.lean
+++ b/src/Lean/Meta/Sym/Simp/DiscrTree.lean
@@ -167,9 +167,18 @@ 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 k c =>
-    -- Reuse general code path
-    getMatchLoop mctx todo (.node #[] #[(k, c)]) result
+  | .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
@@ -181,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`.

From 5100c1c465cd17fe761325cdc117b904aa118617 Mon Sep 17 00:00:00 2001
From: Rob Simmons <robsimmons@gmail.com>
Date: Wed, 19 Aug 2026 12:05:03 -0400
Subject: [PATCH 4/8] fix mapArraysM behavior

---
 src/Lean/Meta/DiscrTree/Util.lean | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/Lean/Meta/DiscrTree/Util.lean b/src/Lean/Meta/DiscrTree/Util.lean
index 44c64e9cfac0..8a55b12ee0dd 100644
--- a/src/Lean/Meta/DiscrTree/Util.lean
+++ b/src/Lean/Meta/DiscrTree/Util.lean
@@ -173,7 +173,14 @@ 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 => Trie.mapArraysM (.node #[] #[(k, c)]) f
+  | .chain k c => do
+    let vs ← f #[] -- Corner case. TODO: remove this and only call `f` for non-empty arrays in node case?
+    let c ← c.mapArraysM f
+    let cs := if c.isEmptyNode then #[] else #[(k, c)]
+    if vs.size > 0 || c.isEmptyNode then
+      return .node vs cs
+    else
+      return .chain k c
   | .node vs children => do
     let vs ← f vs
     let children ← children.filterMapM fun (k, child) => do

From 8d63dce2433552b5cfa49b6e18b9a32a413b63cb Mon Sep 17 00:00:00 2001
From: Rob Simmons <robsimmons@gmail.com>
Date: Wed, 19 Aug 2026 15:40:24 -0400
Subject: [PATCH 5/8] finish specializing code

---
 src/Lean/Meta/DiscrTree/Main.lean  | 48 +++++++++++++++++++-----------
 src/Lean/Meta/DiscrTree/Types.lean |  5 ++--
 2 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/src/Lean/Meta/DiscrTree/Main.lean b/src/Lean/Meta/DiscrTree/Main.lean
index 8803a288ed18..b34767cabcc9 100644
--- a/src/Lean/Meta/DiscrTree/Main.lean
+++ b/src/Lean/Meta/DiscrTree/Main.lean
@@ -430,7 +430,7 @@ private def getStarResult (d : DiscrTree α) : Array α :=
   let result : Array α := .mkEmpty initCapacity
   match d.root.find? .star with
   | none                  => result
-  | some (.chain _k _c) => panic! "unimpl"
+  | 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 α) :=
@@ -473,7 +473,7 @@ private partial def getMatchLoop (todo : Array Expr) (c : Trie α) (result : Arr
           pure result
       match k with
       | .star  => return result
-      | _  =>
+      | _ =>
         match findKey cs k with
         | none   => return result
         | some c => getMatchLoop (todo ++ args) c.2 result
@@ -558,8 +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 k c =>
-      go (.node #[] #[(k, c)]) result
+    | .chain _ c => go c result
     | .node vs cs =>
       let mut result := result ++ vs
       for (_, trie) in cs do
@@ -593,14 +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
-    | _, .chain k c =>
-      -- Reuse general code path
-      process skip todo (.node #[] #[(k, c)]) result
+    | 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
@@ -609,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 eb0b037177bb..9af3726b6065 100644
--- a/src/Lean/Meta/DiscrTree/Types.lean
+++ b/src/Lean/Meta/DiscrTree/Types.lean
@@ -39,8 +39,9 @@ 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. Special-casing 
+    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 α

From 04ff4e3190f647219cb00d8728a04509dd595fd1 Mon Sep 17 00:00:00 2001
From: Rob Simmons <robsimmons@gmail.com>
Date: Wed, 19 Aug 2026 16:54:05 -0400
Subject: [PATCH 6/8] simplify mapArraysM

---
 src/Lean/Meta/DiscrTree/Util.lean | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/src/Lean/Meta/DiscrTree/Util.lean b/src/Lean/Meta/DiscrTree/Util.lean
index 8a55b12ee0dd..d468221ae335 100644
--- a/src/Lean/Meta/DiscrTree/Util.lean
+++ b/src/Lean/Meta/DiscrTree/Util.lean
@@ -177,10 +177,7 @@ partial def Trie.mapArraysM (t : DiscrTree.Trie α) (f : Array α → m (Array 
     let vs ← f #[] -- Corner case. TODO: remove this and only call `f` for non-empty arrays in node case?
     let c ← c.mapArraysM f
     let cs := if c.isEmptyNode then #[] else #[(k, c)]
-    if vs.size > 0 || c.isEmptyNode then
-      return .node vs cs
-    else
-      return .chain k c
+    return Trie.mkNode vs cs
   | .node vs children => do
     let vs ← f vs
     let children ← children.filterMapM fun (k, child) => do
@@ -189,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]

From 210a648ef2b64980d0839f9dbc30668cd3e4a756 Mon Sep 17 00:00:00 2001
From: Rob Simmons <robsimmons@gmail.com>
Date: Wed, 19 Aug 2026 17:02:30 -0400
Subject: [PATCH 7/8] remove TODOs, fix tests reliant on Trie.format

---
 src/Lean/Meta/DiscrTree/Basic.lean |  5 ++--
 src/Lean/Meta/DiscrTree/Util.lean  |  2 +-
 tests/elab/discrTreeFind.lean      | 36 +++++++++++-----------
 tests/elab/discrTreeOps.lean       | 48 +++++++++++++-----------------
 tests/elab/meta3.lean              |  9 +++---
 5 files changed, 47 insertions(+), 53 deletions(-)

diff --git a/src/Lean/Meta/DiscrTree/Basic.lean b/src/Lean/Meta/DiscrTree/Basic.lean
index 54aa59b10940..e400cb2b0117 100644
--- a/src/Lean/Meta/DiscrTree/Basic.lean
+++ b/src/Lean/Meta/DiscrTree/Basic.lean
@@ -60,9 +60,8 @@ def Key.format : Key → Format
 instance : ToFormat Key := ⟨Key.format⟩
 
 partial def Trie.format [ToFormat α] : Trie α → Format
-  | .chain k c =>
-    -- TODO: Make a separate formatting for chains instead of replicating the node formatting
-    Trie.format (.node #[] #[(k, c)])
+  | .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))
diff --git a/src/Lean/Meta/DiscrTree/Util.lean b/src/Lean/Meta/DiscrTree/Util.lean
index d468221ae335..8e5100aedbb9 100644
--- a/src/Lean/Meta/DiscrTree/Util.lean
+++ b/src/Lean/Meta/DiscrTree/Util.lean
@@ -174,7 +174,7 @@ 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. TODO: remove this and only call `f` for non-empty arrays in node case?
+    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
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]

From 0ad131b78e49e15fc578b21664fc4ea31606c0c8 Mon Sep 17 00:00:00 2001
From: Rob Simmons <robsimmons@gmail.com>
Date: Fri, 28 Aug 2026 13:46:41 -0400
Subject: [PATCH 8/8] Don't panic

---
 src/Lean/Meta/Sym/Simp/DiscrTree.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Lean/Meta/Sym/Simp/DiscrTree.lean b/src/Lean/Meta/Sym/Simp/DiscrTree.lean
index 135be8aa1817..fb59a4181e58 100644
--- a/src/Lean/Meta/Sym/Simp/DiscrTree.lean
+++ b/src/Lean/Meta/Sym/Simp/DiscrTree.lean
@@ -209,7 +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 _ _) => panic! "unimpl"
+  | 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