perf: special case single-child nodes in DiscrTree

src/Lean/Meta/DiscrTree/Basic.lean

@@ -60,6 +60,8 @@ def Key.format : Key → Format
 instance : ToFormat Key := ⟨Key.format⟩
 
 partial def Trie.format [ToFormat α] : Trie α → Format
+  | .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))
@@ -126,7 +128,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] #[]
 
@@ -150,6 +152,16 @@ where
   termination_by vs.size - i
 
 private partial def insertAux [BEq α] (keys : Array Key) (v : α) : Nat → Trie α → Trie α
+  | 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]

src/Lean/Meta/DiscrTree/Main.lean

@@ -430,13 +430,29 @@ private def getStarResult (d : DiscrTree α) : Array α :=
   let result : Array α := .mkEmpty initCapacity
   match d.root.find? .star with
   | none                  => result
+  | 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 α) :=
   cs.binSearch (k, default) (fun a b => a.1 < b.1)
 
 private partial def getMatchLoop (todo : Array Expr) (c : Trie α) (result : Array α) : MetaM (Array α) := do
   match c with
+  | .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
@@ -450,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
@@ -544,6 +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 _ c => go c result
     | .node vs cs =>
       let mut result := result ++ vs
       for (_, trie) in cs do
@@ -577,11 +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
+    | 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
@@ -590,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

src/Lean/Meta/DiscrTree/Types.lean

@@ -38,6 +38,12 @@ 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. 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 α
 
 end DiscrTree

src/Lean/Meta/DiscrTree/Util.lean

@@ -15,6 +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 =>
+    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) =>
@@ -32,6 +34,8 @@ Monadically fold the values stored in a `Trie`.
 -/
 @[specialize]
 partial def foldValuesM [Monad m] (f : σ → α → m σ) : (init : σ) → Trie α → m σ
+  | 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
@@ -47,6 +51,7 @@ def foldValues (f : σ → α → σ) (init : σ) (t : Trie α) : σ :=
 The number of values stored in a `Trie`.
 -/
 partial def size : Trie α → Nat
+  | Trie.chain _ c => size c
   | Trie.node vs children =>
     children.foldl (init := vs.size) fun n (_, c) => n + size c
 
@@ -55,13 +60,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⟩
 
 /--
@@ -70,6 +79,7 @@ Equivalent to `t.asNode.1`.
 -/
 @[inline]
 def nodeValues : Trie α → Array α
+  | .chain _ _ => #[]
   | .node vs _ => vs
 
 /--
@@ -78,6 +88,7 @@ Equivalent to `t.asNode.2`.
 -/
 @[inline]
 def nodeChildren : Trie α → Array (Key × Trie α)
+  | .chain k v => #[(k, v)]
   | .node _ cs => cs
 
 /--
@@ -88,6 +99,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
@@ -161,6 +173,11 @@ 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 => do
+    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
   | .node vs children => do
     let vs ← f vs
     let children ← children.filterMapM fun (k, child) => do
@@ -169,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]

src/Lean/Meta/Sym/Simp/DiscrTree.lean

@@ -167,6 +167,19 @@ 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 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
     if todo.isEmpty then
@@ -177,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`.
@@ -204,6 +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 _ _) => .mkEmpty initCapacity -- unreachable in well-formed trees!
   | some (.node vs _) => vs
   let e := resolveAssignedMVars mctx <| etaReduce e
   match d.root.find? (getKey e) with