refactor: path compression for DiscrTree

This PR adds path compression to the DiscrTree trie data structure.
Instead of a single `.node` constructor, the `Trie` type now has four
constructors: `.empty`, `.values`, `.path` (for compressed sequences of
keys with no branching), and `.branch`. This reduces the size of large
discrimination trees (e.g. mathlib's library_search DiscrTree from
217 MB to 41 MB) and the maximum depth (from 2596 to 27).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: nomeata

src/Lean/Meta/DiscrTree/Basic.lean

@@ -19,7 +19,7 @@ def hasNoindexAnnotation (e : Expr) : Bool :=
   annotation? `noindex e |>.isSome
 
 instance : Inhabited (Trie α) where
-  default := .node #[] #[]
+  default := .empty
 
 instance : Inhabited (DiscrTree α) where
   default := {}
@@ -60,9 +60,17 @@ def Key.format : Key → Format
 instance : ToFormat Key := ⟨Key.format⟩
 
 partial def Trie.format [ToFormat α] : Trie α → Format
-  | .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))
+  | .empty => "empty"
+  | .values vs t => Format.group $ Format.paren $
+    "values" ++ (if vs.isEmpty then Format.nil else " " ++ Std.format vs)
+    ++ Format.line ++ format t
+  | .path ks t => Format.group $ Format.paren $
+    "path" ++ (if ks.isEmpty then Format.nil else " " ++ Std.format ks)
+    ++ Format.line ++ format t
+  | .branch cs => Format.group $ Format.paren $
+    "node"
+    ++ Format.join (cs.toList.map fun ⟨k, c⟩ =>
+      Format.line ++ Format.paren (Std.format k ++ " => " ++ format c))
 
 instance [ToFormat α] : ToFormat (Trie α) := ⟨Trie.format⟩
 
@@ -122,13 +130,13 @@ where
       r := r.push (← go)
     return r
 
+/-- Creates a trie with the keys `keys` starting at `i`, and the value `v` as the leaf -/
 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)]
+  let t := .values #[v] .empty
+  if i < keys.size then
+    .path (keys.extract i keys.size) t
   else
-    .node #[v] #[]
+    t
 
 /--
 If `vs` contains an element `v'` such that `v == v'`, then replace `v'` with `v`.
@@ -149,18 +157,78 @@ where
       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
+/--
+Calculate the length of the common prefix of two arrays of keys.
+The parameter `i` marks the starting position in the second array.
+-/
+def commonPrefix (ks1 : Array Key) (ks2 : Array Key) (i : Nat) : Nat :=
+  go 0
+where
+  go (j : Nat) : Nat :=
+    if h1 : j < ks1.size then
+      if h2 : j + i < ks2.size then
+        if ks1[j] == ks2[j + i] then
+          go (j + 1)
+        else
+          j
+      else
+        j
+    else
+      j
+  termination_by ks1.size - j
+
+/-- Smart constructor around branch that ensures the ordering -/
+private def branch2 (k1 : Key) (t1 : Trie α) (k2 : Key) (t2 : Trie α) : Trie α :=
+  if k1 < k2 then
+    .branch #[(k1, t1), (k2, t2)]
+  else
+    .branch #[(k2, t2), (k1, t1)]
+
+/-- Smart constructor ensuring that `.values` constructors are not nested -/
+private partial def insertHere [BEq α] (v : α) : Trie α → Trie α
+  | .values vs t => .values (insertVal vs v) t
+  | t => .values #[v] t
+
+private partial def insertAux [BEq α] (keys : Array Key) (v : α) (i : Nat) (t : Trie α) :
+    Trie α :=
+  if h : i < keys.size then
+    -- we have to walk down the tree some more
+    match t with
+    | .empty => createNodes keys v i
+    | .values _ t => insertAux keys v i t
+    | .path ks t =>
+      let j := commonPrefix ks keys i
+      let t' := -- the new tree after the common prefix
+        if h1 : j < ks.size then
+          if h2 : i + j < keys.size then
+            -- we must branch at offset j
+            let k1 := ks[j]
+            let t1 := if j + 1 < ks.size then .path (ks.extract (j + 1) ks.size) t else t
+            let k2 := keys[i + j]
+            let t2 := createNodes keys v (i + j + 1)
+            branch2 k1 t1 k2 t2
+          else
+            -- the entry keys are a prefix of the path in the node: split the path, insert the value
+            .values #[v] (.path (ks.extract j ks.size) t)
+        else
+          -- the node path is a prefix of the new entry
+          insertAux keys v (i + j) t
+      if 0 < j then
+        -- add a .path for the common prefix, if present
+        .path (ks.extract 0 j) t'
+      else
+        t'
+    | .branch cs =>
       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
+      .branch c
+  else
+    -- this is where we need to insert the value
+    insertHere v t
 
 def insertKeyValue [BEq α] (d : DiscrTree α) (keys : Array Key) (v : α) : DiscrTree α :=
   if keys.isEmpty then panic! "invalid key sequence"
@@ -174,6 +242,10 @@ def insertKeyValue [BEq α] (d : DiscrTree α) (keys : Array Key) (v : α) : Dis
       let c := insertAux keys v 1 c
       { root := d.root.insert k c }
 
+def getValues : Trie α → Array α
+  | .values vs _ => vs
+  | _ => #[]
+
 @[deprecated insertKeyValue (since := "2026-01-02")]
 def insertCore [BEq α] (d : DiscrTree α) (keys : Array Key) (v : α) : DiscrTree α :=
   insertKeyValue d keys v

src/Lean/Meta/DiscrTree/Main.lean

@@ -429,27 +429,27 @@ private abbrev getUnifyKeyArgs (e : Expr) (root : Bool) : MetaM (Key × Array Ex
 private def getStarResult (d : DiscrTree α) : Array α :=
   let result : Array α := .mkEmpty initCapacity
   match d.root.find? .star with
-  | none                  => result
-  | some (.node vs _) => result ++ vs
+  | none   => result
+  | some t => result ++ getValues t
 
 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
-  | .node vs cs =>
-    if todo.isEmpty then
-      return result ++ vs
-    else if cs.isEmpty then
-      return result
-    else
+  if todo.isEmpty then
+    return result ++ getValues c
+  else
+    match c with
+    | .empty => return result
+    | .values _ t => getMatchLoop todo t result
+    | .branch cs =>
+      if cs.isEmpty then return result else -- should not happen
       let e     := todo.back!
       let todo  := todo.pop
       let first := cs[0]! /- Recall that `Key.star` is the minimal key -/
-      let (k, args) ← getMatchKeyArgs e (root := false)
       /- 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`. -/
+        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 α) :=
         if first.1 == .star then
           getMatchLoop todo first.2 result
@@ -459,10 +459,38 @@ private partial def getMatchLoop (todo : Array Expr) (c : Trie α) (result : Arr
         match findKey cs k with
         | none   => return result
         | some c => getMatchLoop (todo ++ args) c.2 result
+      let (k, args) ← getMatchKeyArgs e (root := false)
       let result ← visitStar result
       match k with
       | .star  => return result
       | _      => visitNonStar k args result
+    | .path ks t =>
+      let rec loop (todo : Array Expr) (result : Array α) (i : Nat) : MetaM (Array α) := do
+        -- the following logic is a copy of the .branch case, as if `cs` is a singleton
+        if h : i < ks.size then
+          if todo.isEmpty then
+            return result
+          let e     := todo.back!
+          let todo  := todo.pop
+          let k'    := ks[i]
+          let visitStar (result : Array α) : MetaM (Array α) :=
+            if k' == .star then
+              loop todo result (i + 1)
+            else
+              return result
+          let visitNonStar (k : Key) (args : Array Expr) (result : Array α) : MetaM (Array α) :=
+            if k == k' then
+              loop (todo ++ args) result (i + 1)
+            else
+              return result
+          let (k, args) ← getMatchKeyArgs e (root := false)
+          let result ← visitStar result
+          match k with
+            | .star => return result
+            | _     => visitNonStar k args result
+        else
+          getMatchLoop todo t result
+      loop todo result 0
 
 private def getMatchRoot (d : DiscrTree α) (k : Key) (args : Array Expr) (result : Array α) : MetaM (Array α) :=
   match d.root.find? k with
@@ -544,8 +572,11 @@ private partial def getAllValuesForKey (d : DiscrTree α) (k : Key) (result : Ar
 where
   go (trie : Trie α) (result : Array α) : Array α := Id.run do
     match trie with
-    | .node vs cs =>
-      let mut result := result ++ vs
+    | .empty => return result
+    | .values vs t => go t (result ++ vs)
+    | .path _ t => go t result
+    | .branch cs =>
+      let mut result := result
       for (_, trie) in cs do
         result := go trie result
       return result
@@ -576,33 +607,66 @@ partial def getUnify (d : DiscrTree α) (e : Expr) : MetaM (Array α) :=
       | some c => process 0 args c result
 where
   process (skip : Nat) (todo : Array Expr) (c : Trie α) (result : Array α) : MetaM (Array α) := do
-    match skip, c with
-    | 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, .node vs cs => do
-      if todo.isEmpty then
-        return result ++ vs
-      else if cs.isEmpty then
-        return result
-      else
-        let e     := todo.back!
-        let todo  := todo.pop
-        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
+    if skip == 0 && todo.isEmpty then
+      return result ++ getValues c
+    else match c with
+      | .empty => return result
+      | .values _ t => process skip todo t result
+      | .branch cs =>
+        match skip with
+        | skip+1 =>
+          if cs.isEmpty then
             return result
-        let visitNonStar (k : Key) (args : Array Expr) (result : Array α) : MetaM (Array α) :=
-          match findKey cs k with
-          | 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)
+          else
+            cs.foldlM (init := result) fun result ⟨k, c⟩ => process (skip + k.arity) todo c result
+        | 0 => do
+          if cs.isEmpty then return result else -- should not happen
+          let e     := todo.back!
+          let todo  := todo.pop
+          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 findKey cs k with
+            | 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)
+      | .path ks t =>
+        let rec loop (skip : Nat) (todo : Array Expr) (result : Array α) (i : Nat) : MetaM (Array α) :=
+          if h : i < ks.size then
+            match skip with
+            | skip+1 =>
+                let k' := ks[i]
+                loop (skip + k'.arity) todo result (i + 1)
+            | 0 => do
+              if todo.isEmpty then
+                return result
+              else
+                let e     := todo.back!
+                let todo  := todo.pop
+                let k'    := ks[i]
+                let visitStar (result : Array α) : MetaM (Array α) :=
+                  if k' == .star then
+                    loop 0 todo result (i + 1)
+                  else
+                    return result
+                let visitNonStar (k : Key) (args : Array Expr) (result : Array α) : MetaM (Array α) :=
+                  if k' == k then
+                    loop 0 (todo ++ args) result (i + 1)
+                  else
+                    return result
+                let (k, args) ← getUnifyKeyArgs e (root := false)
+                match k with
+                | .star  => loop k'.arity todo result (i + 1)
+                | _      => visitNonStar k args (← visitStar result)
+          else
+            process skip todo t result
+        loop skip todo result 0
 
 end Lean.Meta.DiscrTree

src/Lean/Meta/DiscrTree/Types.lean

@@ -38,7 +38,10 @@ instance : Hashable Key := ⟨Key.hash⟩
 Discrimination tree trie. See `DiscrTree`.
 -/
 inductive Trie (α : Type) where
-  | node (vs : Array α) (children : Array (Key × Trie α)) : Trie α
+  | empty
+  | values (vs : Array α) (t : Trie α)
+  | path (ks : Array Key) (t : Trie α)
+  | branch (children : Array (Key × Trie α))
 
 end DiscrTree