chore:

src/Lean/Compiler/LCNF/Main.lean

@@ -4,16 +4,14 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Compiler.Options
 public import Lean.Compiler.IR
 public import Lean.Compiler.LCNF.Passes
 public import Lean.Compiler.LCNF.ToDecl
 public import Lean.Compiler.LCNF.Check
-
+import Lean.Meta.Match.MatcherInfo
 public section
-
 namespace Lean.Compiler.LCNF
 /--
 We do not generate code for `declName` if

src/Lean/Compiler/LCNF/Specialize.lean

@@ -4,7 +4,6 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Compiler.LCNF.SpecInfo
 public import Lean.Compiler.LCNF.MonadScope
@@ -13,7 +12,7 @@ import Lean.Compiler.LCNF.Simp
 import Lean.Compiler.LCNF.ToExpr
 import Lean.Compiler.LCNF.Level
 import Lean.Compiler.LCNF.Closure
-
+import Lean.Meta.Transform
 namespace Lean.Compiler.LCNF
 namespace Specialize
 

src/Lean/Compiler/LCNF/ToDecl.lean

@@ -4,13 +4,12 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Compiler.InitAttr
 public import Lean.Compiler.LCNF.ToLCNF
-
+import Lean.Meta.Transform
+import Lean.Meta.Match.MatcherInfo
 public section
-
 namespace Lean.Compiler.LCNF
 /--
 Inline constants tagged with the `[macro_inline]` attribute occurring in `e`.

src/Lean/Compiler/LCNF/ToLCNF.lean

@@ -4,7 +4,6 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Meta.AppBuilder
 public import Lean.Compiler.CSimpAttr
@@ -12,9 +11,8 @@ public import Lean.Compiler.ImplementedByAttr
 public import Lean.Compiler.LCNF.Bind
 public import Lean.Compiler.NeverExtractAttr
 import Lean.Meta.CasesInfo
-
+import Lean.Meta.WHNF
 public section
-
 namespace Lean.Compiler.LCNF
 namespace ToLCNF
 

src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/ReifiedBVExpr.lean

@@ -4,13 +4,11 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Henrik Böving
 -/
 module
-
 prelude
 public import Lean.Elab.Tactic.BVDecide.Frontend.BVDecide.Reflect
 public import Std.Tactic.BVDecide.Reflect
-
+import Lean.Meta.LitValues
 public section
-
 /-!
 Provides the logic for reifying `BitVec` expressions.
 -/

src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/Reify.lean

@@ -4,10 +4,9 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Henrik Böving
 -/
 module
-
 prelude
 public import Lean.Elab.Tactic.BVDecide.Frontend.BVDecide.ReifiedLemmas
-
+import Lean.Meta.LitValues
 public section
 
 /-!

src/Lean/Elab/Tactic/Injection.lean

@@ -4,12 +4,10 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Meta.Tactic.Injection
 public import Lean.Meta.Tactic.Assumption
 public import Lean.Elab.Tactic.ElabTerm
-
 public section
 namespace Lean.Elab.Tactic
 

src/Lean/Meta/ACLt.lean

@@ -4,12 +4,10 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-public import Lean.Meta.DiscrTree
-
+public import Lean.Meta.DiscrTree.Main
+import Lean.Meta.WHNF
 public section
-
 namespace Lean
 
 def Expr.ctorWeight : Expr → UInt8

src/Lean/Meta/AppBuilder.lean

@@ -4,15 +4,14 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Meta.SynthInstance
 public import Lean.Meta.DecLevel
 import Lean.Meta.SameCtorUtils
 import Lean.Data.Array
-
+import Lean.Meta.CtorRecognizer
+import Lean.Structure
 public section
-
 namespace Lean.Meta
 
 /-- Returns `id e` -/

src/Lean/Meta/Coe.lean

@@ -4,15 +4,14 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Meta.AppBuilder
 import Lean.ExtraModUses
-
+import Lean.ProjFns
+import Lean.Meta.Transform
+import Lean.Meta.WHNF
 public section
-
 namespace Lean.Meta
-
 /--
 Tags declarations to be unfolded during coercion elaboration.
 

src/Lean/Meta/CongrTheorems.lean

@@ -4,14 +4,13 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.AddDecl
 public import Lean.ReservedNameAction
+import Lean.Structure
 import Lean.Meta.Tactic.Subst
-
+import Lean.Meta.FunInfo
 public section
-
 namespace Lean.Meta
 
 inductive CongrArgKind where

src/Lean/Meta/Constructions/SparseCasesOn.lean

@@ -3,16 +3,14 @@ Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Joachim Breitner
 -/
-
 module
-
 prelude
 public import Lean.Meta.Basic
 import Lean.AddDecl
 import Lean.Meta.Constructions.CtorIdx
 import Lean.Meta.AppBuilder
 import Lean.Meta.HasNotBit
-
+import Lean.Meta.WHNF
 /-!  See `mkSparseCasesOn` below.  -/
 
 namespace Lean.Meta

src/Lean/Meta/DiscrTree.lean

@@ -4,874 +4,8 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura, Jannis Limperg, Kim Morrison
 -/
 module
-
 prelude
 public import Lean.Meta.WHNF
-public import Lean.Meta.DiscrTreeTypes
-
-public section
-
-namespace Lean.Meta.DiscrTree
-/-!
-  (Imperfect) discrimination trees.
-  We use a hybrid representation.
-  - A `PersistentHashMap` for the root node which usually contains many children.
-  - A sorted array of key/node pairs for inner nodes.
-
-  The edges are labeled by keys:
-  - Constant names (and arity). Universe levels are ignored.
-  - Free variables (and arity). Thus, an entry in the discrimination tree
-    may reference hypotheses from the local context.
-  - Literals
-  - Star/Wildcard. We use them to represent metavariables and terms
-    we want to ignore. We ignore implicit arguments and proofs.
-  - Other. We use to represent other kinds of terms (e.g., nested lambda, forall, sort, etc).
-
-  We reduce terms using `TransparencyMode.reducible`. Thus, all reducible
-  definitions in an expression `e` are unfolded before we insert it into the
-  discrimination tree.
-
-  Recall that projections from classes are **NOT** reducible.
-  For example, the expressions `Add.add α (ringAdd ?α ?s) ?x ?x`
-  and `Add.add Nat Nat.hasAdd a b` generates paths with the following keys
-  respectively
-  ```
-  ⟨Add.add, 4⟩, α, *, *, *
-  ⟨Add.add, 4⟩, Nat, *, ⟨a,0⟩, ⟨b,0⟩
-  ```
-
-  That is, we don't reduce `Add.add Nat inst a b` into `Nat.add a b`.
-  We say the `Add.add` applications are the de-facto canonical forms in
-  the metaprogramming framework.
-  Moreover, it is the metaprogrammer's responsibility to re-pack applications such as
-  `Nat.add a b` into `Add.add Nat inst a b`.
-
-  Remark: we store the arity in the keys
-  1- To be able to implement the "skip" operation when retrieving "candidate"
-     unifiers.
-  2- Distinguish partial applications `f a`, `f a b`, and `f a b c`.
--/
-
-def Key.lt : Key → Key → Bool
-  | .lit v₁,        .lit v₂        => v₁ < v₂
-  | .fvar n₁ a₁,    .fvar n₂ a₂    => Name.quickLt n₁.name n₂.name || (n₁ == n₂ && a₁ < a₂)
-  | .const n₁ a₁,   .const n₂ a₂   => Name.quickLt n₁ n₂ || (n₁ == n₂ && a₁ < a₂)
-  | .proj s₁ i₁ a₁, .proj s₂ i₂ a₂ => Name.quickLt s₁ s₂ || (s₁ == s₂ && i₁ < i₂) || (s₁ == s₂ && i₁ == i₂ && a₁ < a₂)
-  | k₁,             k₂             => k₁.ctorIdx < k₂.ctorIdx
-
-instance : LT Key := ⟨fun a b => Key.lt a b⟩
-instance (a b : Key) : Decidable (a < b) := inferInstanceAs (Decidable (Key.lt a b))
-
-def Key.format : Key → Format
-  | .star            => "*"
-  | .other           => "◾"
-  | .lit (.natVal v) => Std.format v
-  | .lit (.strVal v) => repr v
-  | .const k _       => Std.format k
-  | .proj s i _      => Std.format s ++ "." ++ Std.format i
-  | .fvar k _        => Std.format k.name
-  | .arrow           => "∀"
-
-instance : ToFormat Key := ⟨Key.format⟩
-
-/--
-Helper function for converting an entry (i.e., `Array Key`) to the discrimination tree into
-`MessageData` that is more user-friendly. We use this function to implement diagnostic information.
--/
-partial def keysAsPattern (keys : Array Key) : CoreM MessageData := do
-  go (parenIfNonAtomic := false) |>.run' keys.toList
-where
-  next? : StateRefT (List Key) CoreM (Option Key) := do
-    let key :: keys ← get | return none
-    set keys
-    return some key
-
-  mkApp (f : MessageData) (args : Array MessageData) (parenIfNonAtomic : Bool) : CoreM MessageData := do
-    if args.isEmpty then
-      return f
-    else
-      let mut r := m!""
-      for arg in args do
-        r := r ++ Format.line ++ arg
-      r := f ++ .nest 2 r
-      if parenIfNonAtomic then
-        return .paren r
-      else
-        return .group r
-
-  go (parenIfNonAtomic := true) : StateRefT (List Key) CoreM MessageData := do
-    let some key ← next? | return .nil
-    match key with
-    | .const declName nargs =>
-      mkApp m!"{← mkConstWithLevelParams declName}" (← goN nargs) parenIfNonAtomic
-    | .fvar fvarId nargs =>
-      mkApp m!"{mkFVar fvarId}" (← goN nargs) parenIfNonAtomic
-    | .proj _ i nargs =>
-      mkApp m!"{← go}.{i+1}" (← goN nargs) parenIfNonAtomic
-    | .arrow =>
-      mkApp m!"∀ " (← goN 1) parenIfNonAtomic
-    | .star => return "_"
-    | .other => return "<other>"
-    | .lit (.natVal v) => return m!"{v}"
-    | .lit (.strVal v) => return m!"{v}"
-
-  goN (num : Nat) : StateRefT (List Key) CoreM (Array MessageData) := do
-    let mut r := #[]
-    for _ in *...num do
-      r := r.push (← go)
-    return r
-
-def Key.arity : Key → Nat
-  | .const _ a  => a
-  | .fvar _ a   => a
-  /-
-  Remark: `.arrow` used to have arity 2, and was used to encode only **non**-dependent
-  arrows. However, this feature was a recurrent source of bugs. For example, a
-  theorem about a dependent arrow can be applied to a non-dependent one. The
-  reverse direction may also happen. See issue #2835. Therefore, `.arrow` was made
-  to have arity 0. But this throws away easy to use information, and makes it so
-  that ∀ and ∃ behave quite differently. So now `.arrow` at least indexes the
-  domain of the forall (whether dependent or non-dependent).
-  -/
-  | .arrow      => 1
-  | .proj _ _ a => 1 + a
-  | _           => 0
-
-instance : Inhabited (Trie α) := ⟨.node #[] #[]⟩
-
-def empty : DiscrTree α := { root := {} }
-
-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))
-
-instance [ToFormat α] : ToFormat (Trie α) := ⟨Trie.format⟩
-
-partial def format [ToFormat α] (d : DiscrTree α) : Format :=
-  let (_, r) := d.root.foldl
-    (fun (p : Bool × Format) k c =>
-      (false, p.2 ++ (if p.1 then Format.nil else Format.line) ++ Format.paren (Std.format k ++ " => " ++ Std.format c)))
-    (true, Format.nil)
-  Format.group r
-
-instance [ToFormat α] : ToFormat (DiscrTree α) := ⟨format⟩
-
-/-- The discrimination tree ignores implicit arguments and proofs.
-   We use the following auxiliary id as a "mark". -/
-private def tmpMVarId : MVarId := { name := `_discr_tree_tmp }
-private def tmpStar := mkMVar tmpMVarId
-
-instance : Inhabited (DiscrTree α) where
-  default := {}
-
-/--
-  Return true iff the argument should be treated as a "wildcard" by the discrimination tree.
-
-  - We ignore proofs because of proof irrelevance. It doesn't make sense to try to
-    index their structure.
-
-  - We ignore instance implicit arguments (e.g., `[Add α]`) because they are "morally" canonical.
-    Moreover, we may have many definitionally equal terms floating around.
-    Example: `Ring.hasAdd Int Int.isRing` and `Int.hasAdd`.
-
-  - We considered ignoring implicit arguments (e.g., `{α : Type}`) since users don't "see" them,
-    and may not even understand why some simplification rule is not firing.
-    However, in type class resolution, we have instance such as `Decidable (@Eq Nat x y)`,
-    where `Nat` is an implicit argument. Thus, we would add the path
-    ```
-    Decidable -> Eq -> * -> * -> * -> [Nat.decEq]
-    ```
-    to the discrimination tree IF we ignored the implicit `Nat` argument.
-    This would be BAD since **ALL** decidable equality instances would be in the same path.
-    So, we index implicit arguments if they are types.
-    This setting seems sensible for simplification theorems such as:
-    ```
-    forall (x y : Unit), (@Eq Unit x y) = true
-    ```
-    If we ignore the implicit argument `Unit`, the `DiscrTree` will say it is a candidate
-    simplification theorem for any equality in our goal.
-
-  Remark: if users have problems with the solution above, we may provide a `noIndexing` annotation,
-  and `ignoreArg` would return true for any term of the form `noIndexing t`.
--/
-private def ignoreArg (a : Expr) (i : Nat) (infos : Array ParamInfo) : MetaM Bool := do
-  if h : i < infos.size then
-    let info := infos[i]
-    if info.isInstImplicit then
-      return true
-    else if info.isImplicit || info.isStrictImplicit then
-      return !(← isType a)
-    else
-      isProof a
-  else
-    isProof a
-
-private partial def pushArgsAux (infos : Array ParamInfo) : Nat → Expr → Array Expr → MetaM (Array Expr)
-  | i, .app f a, todo => do
-    if (← ignoreArg a i infos) then
-      pushArgsAux infos (i-1) f (todo.push tmpStar)
-    else
-      pushArgsAux infos (i-1) f (todo.push a)
-  | _, _, todo => return todo
-
-/--
-  Return true if `e` is one of the following
-  - A nat literal (numeral)
-  - `Nat.zero`
-  - `Nat.succ x` where `isNumeral x`
-  - `OfNat.ofNat _ x _` where `isNumeral x` -/
-private partial def isNumeral (e : Expr) : Bool :=
-  if e.isRawNatLit then true
-  else
-    let f := e.getAppFn
-    if !f.isConst then false
-    else
-      let fName := f.constName!
-      if fName == ``Nat.succ && e.getAppNumArgs == 1 then isNumeral e.appArg!
-      else if fName == ``OfNat.ofNat && e.getAppNumArgs == 3 then isNumeral (e.getArg! 1)
-      else if fName == ``Nat.zero && e.getAppNumArgs == 0 then true
-      else false
-
-private partial def toNatLit? (e : Expr) : Option Literal :=
-  if isNumeral e then
-    if let some n := loop e then
-      some (.natVal n)
-    else
-      none
-  else
-    none
-where
-  loop (e : Expr) : OptionT Id Nat := do
-    let f := e.getAppFn
-    match f with
-    | .lit (.natVal n) => return n
-    | .const fName .. =>
-      if fName == ``Nat.succ && e.getAppNumArgs == 1 then
-        let r ← loop e.appArg!
-        return r+1
-      else if fName == ``OfNat.ofNat && e.getAppNumArgs == 3 then
-        loop (e.getArg! 1)
-      else if fName == ``Nat.zero && e.getAppNumArgs == 0 then
-        return 0
-      else
-        failure
-    | _ => failure
-
-private def isNatType (e : Expr) : MetaM Bool :=
-  return (← whnf e).isConstOf ``Nat
-
-/--
-  Return true if `e` is one of the following
-  - `Nat.add _ k` where `isNumeral k`
-  - `Add.add Nat _ _ k` where `isNumeral k`
-  - `HAdd.hAdd _ Nat _ _ k` where `isNumeral k`
-  - `Nat.succ _`
-  This function assumes `e.isAppOf fName`
--/
-private def isOffset (fName : Name) (e : Expr) : MetaM Bool := do
-  if fName == ``Nat.add && e.getAppNumArgs == 2 then
-    return isNumeral e.appArg!
-  else if fName == ``Add.add && e.getAppNumArgs == 4 then
-    if (← isNatType (e.getArg! 0)) then return isNumeral e.appArg! else return false
-  else if fName == ``HAdd.hAdd && e.getAppNumArgs == 6 then
-    if (← isNatType (e.getArg! 1)) then return isNumeral e.appArg! else return false
-  else
-    return fName == ``Nat.succ && e.getAppNumArgs == 1
-
-/--
-  TODO: add hook for users adding their own functions for controlling `shouldAddAsStar`
-  Different `DiscrTree` users may populate this set using, for example, attributes.
-
-  Remark: we currently tag "offset" terms as star to avoid having to add special
-  support for offset terms.
-  Example, suppose the discrimination tree contains the entry
-  `Nat.succ ?m |-> v`, and we are trying to retrieve the matches for `Expr.lit (Literal.natVal 1) _`.
-  In this scenario, we want to retrieve `Nat.succ ?m |-> v`
--/
-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.
--/
-partial def reduce (e : Expr) : MetaM Expr := do
-  let e ← whnfCore e
-  match (← unfoldDefinition? e) with
-  | some e => reduce e
-  | none => match e.etaExpandedStrict? with
-    | some e => reduce e
-    | none   => return e
-
-/--
-  Return `true` if `fn` is a "bad" key. That is, `pushArgs` would add `Key.other` or `Key.star`.
-  We use this function when processing "root terms, and will avoid unfolding terms.
-  Note that without this trick the pattern `List.map f ∘ List.map g` would be mapped into the key `Key.other`
-  since the function composition `∘` would be unfolded and we would get `fun x => List.map g (List.map f x)`
--/
-private def isBadKey (fn : Expr) : Bool :=
-  match fn with
-  | .lit ..     => false
-  | .const ..   => false
-  | .fvar ..    => false
-  | .proj ..    => false
-  | .forallE .. => false
-  | _           => true
-
-/--
-Reduce `e` until we get an irreducible term (modulo current reducibility setting) or the resulting term
-is a bad key (see comment at `isBadKey`).
-We use this method instead of `reduce` for root terms at `pushArgs`. -/
-private partial def reduceUntilBadKey (e : Expr) : MetaM Expr := do
-  let e ← step e
-  match e.etaExpandedStrict? with
-  | some e => reduceUntilBadKey e
-  | none   => return e
-where
-  step (e : Expr) := do
-    let e ← whnfCore e
-    match (← unfoldDefinition? e) with
-    | some e' => if isBadKey e'.getAppFn then return e else step e'
-    | none    => return e
-
-/-- whnf for the discrimination tree module -/
-def reduceDT (e : Expr) (root : Bool) : MetaM Expr :=
-  if root then reduceUntilBadKey e else reduce e
-
-/- Remark: we use `shouldAddAsStar` only for nested terms, and `root == false` for nested terms -/
-
-/--
-Append `n` wildcards to `todo`
--/
-private def pushWildcards (n : Nat) (todo : Array Expr) : Array Expr :=
-  match n with
-  | 0   => todo
-  | n+1 => pushWildcards n (todo.push tmpStar)
-
-/--
-When `noIndexAtArgs := true`, `pushArgs` assumes function application arguments have a `no_index` annotation.
-That is, `f a b` is indexed as it was `f (no_index a) (no_index b)`.
-This feature is used when indexing local proofs in the simplifier. This is useful in examples like the one described on issue #2670.
-In this issue, we have a local hypotheses `(h : ∀ p : α × β, f p p.2 = p.2)`, and users expect it to be applicable to
-`f (a, b) b = b`. This worked in Lean 3 since no indexing was used. We can retrieve Lean 3 behavior by writing
-`(h : ∀ p : α × β, f p (no_index p.2) = p.2)`, but this is very inconvenient when the hypotheses was not written by the user in first place.
-For example, it was introduced by another tactic. Thus, when populating the discrimination tree explicit arguments provided to `simp` (e.g., `simp [h]`),
-we use `noIndexAtArgs := true`. See comment: https://github.com/leanprover/lean4/issues/2670#issuecomment-1758889365
--/
-private def pushArgs (root : Bool) (todo : Array Expr) (e : Expr) (noIndexAtArgs : Bool) : MetaM (Key × Array Expr) := do
-  if hasNoindexAnnotation e then
-    return (.star, todo)
-  else
-    let e ← reduceDT e root
-    let fn := e.getAppFn
-    let push (k : Key) (nargs : Nat) (todo : Array Expr): MetaM (Key × Array Expr) := do
-      let info ← getFunInfoNArgs fn nargs
-      let todo ← if noIndexAtArgs then
-        pure <| pushWildcards nargs todo
-      else
-        pushArgsAux info.paramInfo (nargs-1) e todo
-      return (k, todo)
-    match fn with
-    | .lit v     =>
-      return (.lit v, todo)
-    | .const c _ =>
-      unless root do
-        if let some v := toNatLit? e then
-          return (.lit v, todo)
-        if (← shouldAddAsStar c e) then
-          return (.star, todo)
-      let nargs := e.getAppNumArgs
-      push (.const c nargs) nargs todo
-    | .proj s i a =>
-      /-
-      If `s` is a class, then `a` is an instance. Thus, we annotate `a` with `no_index` since we do not
-      index instances. This should only happen if users mark a class projection function as `[reducible]`.
-
-      TODO: add better support for projections that are functions
-      -/
-      let a := if isClass (← getEnv) s then mkNoindexAnnotation a else a
-      let nargs := e.getAppNumArgs
-      push (.proj s i nargs) nargs (todo.push a)
-    | .fvar fvarId   =>
-      let nargs := e.getAppNumArgs
-      push (.fvar fvarId nargs) nargs todo
-    | .mvar mvarId   =>
-      if mvarId == tmpMVarId then
-        -- We use `tmp to mark implicit arguments and proofs
-        return (.star, todo)
-      else if (← mvarId.isReadOnlyOrSyntheticOpaque) then
-        return (.other, todo)
-      else
-        return (.star, todo)
-    | .forallE _n d _ _ =>
-      return (.arrow, todo.push d)
-    | _ => return (.other, todo)
-
-@[inherit_doc pushArgs]
-partial def mkPathAux (root : Bool) (todo : Array Expr) (keys : Array Key) (noIndexAtArgs : Bool) : MetaM (Array Key) := do
-  if todo.isEmpty then
-    return keys
-  else
-    let e    := todo.back!
-    let todo := todo.pop
-    let (k, todo) ← pushArgs root todo e noIndexAtArgs
-    mkPathAux false todo (keys.push k) noIndexAtArgs
-
-private def initCapacity := 8
-
-@[inherit_doc pushArgs]
-def mkPath (e : Expr) (noIndexAtArgs := false) : MetaM (Array Key) := do
-  withReducible do
-    let todo : Array Expr := .mkEmpty initCapacity
-    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
-
-/--
-Inserts a value into a discrimination tree,
-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
-  else
-    d.insertCore keys v
-
-private def getKeyArgs (e : Expr) (isMatch root : Bool) : MetaM (Key × Array Expr) := do
-  let e ← reduceDT e root
-  unless root do
-    -- See pushArgs
-    if let some v := toNatLit? e then
-      return (.lit v, #[])
-  match e.getAppFn with
-  | .lit v         => return (.lit v, #[])
-  | .const c _     =>
-    if (← getConfig).isDefEqStuckEx && e.hasExprMVar then
-      if (← isReducible c) then
-        /- `e` is a term `c ...` s.t. `c` is reducible and `e` has metavariables, but it was not unfolded.
-           This can happen if the metavariables in `e` are "blocking" smart unfolding.
-           If `isDefEqStuckEx` is enabled, then we must throw the `isDefEqStuck` exception to postpone TC resolution.
-           Here is an example. Suppose we have
-           ```
-            inductive Ty where
-              | bool | fn (a ty : Ty)
-
-
-            @[reducible] def Ty.interp : Ty → Type
-              | bool   => Bool
-              | fn a b => a.interp → b.interp
-           ```
-           and we are trying to synthesize `BEq (Ty.interp ?m)`
-        -/
-        Meta.throwIsDefEqStuck
-      else if let some matcherInfo := isMatcherAppCore? (← getEnv) e then
-        -- A matcher application is stuck is one of the discriminants has a metavariable
-        let args := e.getAppArgs
-        for arg in args[matcherInfo.getFirstDiscrPos...(matcherInfo.getFirstDiscrPos + matcherInfo.numDiscrs)] do
-          if arg.hasExprMVar then
-            Meta.throwIsDefEqStuck
-      else if (← isRec c) then
-        /- Similar to the previous case, but for `match` and recursor applications. It may be stuck (i.e., did not reduce)
-           because of metavariables. -/
-        Meta.throwIsDefEqStuck
-    let nargs := e.getAppNumArgs
-    return (.const c nargs, e.getAppRevArgs)
-  | .fvar fvarId   =>
-    let nargs := e.getAppNumArgs
-    return (.fvar fvarId nargs, e.getAppRevArgs)
-  | .mvar mvarId   =>
-    if isMatch then
-      return (.other, #[])
-    else do
-      let cfg ← getConfig
-      if cfg.isDefEqStuckEx then
-        /-
-          When the configuration flag `isDefEqStuckEx` is set to true,
-          we want `isDefEq` to throw an exception whenever it tries to assign
-          a read-only metavariable.
-          This feature is useful for type class resolution where
-          we may want to notify the caller that the TC problem may be solvable
-          later after it assigns `?m`.
-          The method `DiscrTree.getUnify e` returns candidates `c` that may "unify" with `e`.
-          That is, `isDefEq c e` may return true. Now, consider `DiscrTree.getUnify d (Add ?m)`
-          where `?m` is a read-only metavariable, and the discrimination tree contains the keys
-          `HadAdd Nat` and `Add Int`. If `isDefEqStuckEx` is set to true, we must treat `?m` as
-          a regular metavariable here, otherwise we return the empty set of candidates.
-          This is incorrect because it is equivalent to saying that there is no solution even if
-          the caller assigns `?m` and try again. -/
-        return (.star, #[])
-      else if (← mvarId.isReadOnlyOrSyntheticOpaque) then
-        return (.other, #[])
-      else
-        return (.star, #[])
-  | .proj s i a .. =>
-    let nargs := e.getAppNumArgs
-    return (.proj s i nargs, #[a] ++ e.getAppRevArgs)
-  | .forallE _ d _ _ => return (.arrow, #[d])
-  | _ => return (.other, #[])
-
-private abbrev getMatchKeyArgs (e : Expr) (root : Bool) : MetaM (Key × Array Expr) :=
-  getKeyArgs e (isMatch := true) (root := root)
-
-private abbrev getUnifyKeyArgs (e : Expr) (root : Bool) : MetaM (Key × Array Expr) :=
-  getKeyArgs e (isMatch := false) (root := root)
-
-private def getStarResult (d : DiscrTree α) : Array α :=
-  let result : Array α := .mkEmpty initCapacity
-  match d.root.find? .star with
-  | none                  => result
-  | 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
-  | .node vs cs =>
-    if todo.isEmpty then
-      return result ++ vs
-    else if cs.isEmpty then
-      return result
-    else
-      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`. -/
-      let visitStar (result : Array α) : MetaM (Array α) :=
-        if first.1 == .star then
-          getMatchLoop 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 => 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
-  | none   => return result
-  | some c => getMatchLoop args c result
-
-private def getMatchCore (d : DiscrTree α) (e : Expr) : MetaM (Key × Array α) :=
-  withReducible do
-    let result := getStarResult d
-    let (k, args) ← getMatchKeyArgs e (root := true)
-    match k with
-    | .star  => return (k, result)
-    | _      => return (k, (← getMatchRoot d k args result))
-
-/--
-  Find values that match `e` in `d`.
--/
-def getMatch (d : DiscrTree α) (e : Expr) : MetaM (Array α) :=
-  return (← getMatchCore d e).2
-
-/--
-  Similar to `getMatch`, but returns solutions that are prefixes of `e`.
-  We store the number of ignored arguments in the result.-/
-partial def getMatchWithExtra (d : DiscrTree α) (e : Expr) : MetaM (Array (α × Nat)) := do
-  let (k, result) ← getMatchCore d e
-  let result := result.map (·, 0)
-  if !e.isApp then
-    return result
-  else if !(← mayMatchPrefix k) then
-    return result
-  else
-    go e.appFn! 1 result
-where
-  mayMatchPrefix (k : Key) : MetaM Bool :=
-    let cont (k : Key) : MetaM Bool :=
-      if d.root.find? k |>.isSome then
-        return true
-      else
-        mayMatchPrefix k
-    match k with
-    | .const f (n+1)  => cont (.const f n)
-    | .fvar f (n+1)   => cont (.fvar f n)
-    | .proj s i (n+1) => cont (.proj s i n)
-    | _               => return false
-
-  go (e : Expr) (numExtra : Nat) (result : Array (α × Nat)) : MetaM (Array (α × Nat)) := do
-    let result := result ++ (← getMatchCore d e).2.map (., numExtra)
-    if e.isApp then
-      go e.appFn! (numExtra + 1) result
-    else
-      return result
-
-/--
-Return the root symbol for `e`, and the number of arguments after `reduceDT`.
--/
-def getMatchKeyRootFor (e : Expr) : MetaM (Key × Nat) := do
-  let e ← reduceDT e (root := true)
-  let numArgs := e.getAppNumArgs
-  let key := match e.getAppFn with
-    | .lit v         => .lit v
-    | .fvar fvarId   => .fvar fvarId numArgs
-    | .mvar _        => .other
-    | .proj s i _ .. => .proj s i numArgs
-    | .forallE ..    => .arrow
-    | .const c _     =>
-      -- This method is used by the simplifier only, we do **not** support
-      -- (← getConfig).isDefEqStuckEx
-      .const c numArgs
-    | _ => .other
-  return (key, numArgs)
-
-/--
-Get all results under key `k`.
--/
-private partial def getAllValuesForKey (d : DiscrTree α) (k : Key) (result : Array α) : Array α :=
-  match d.root.find? k with
-  | none      => result
-  | some trie => go trie result
-where
-  go (trie : Trie α) (result : Array α) : Array α := Id.run do
-    match trie with
-    | .node vs cs =>
-      let mut result := result ++ vs
-      for (_, trie) in cs do
-        result := go trie result
-      return result
-
-/--
-A liberal version of `getMatch` which only takes the root symbol of `e` into account.
-We use this method to simulate Lean 3's indexing.
-
-The natural number in the result is the number of arguments in `e` after `reduceDT`.
--/
-def getMatchLiberal (d : DiscrTree α) (e : Expr) : MetaM (Array α × Nat) := do
-  withReducible do
-    let result := getStarResult d
-    let (k, numArgs) ← getMatchKeyRootFor e
-    match k with
-    | .star  => return (result, numArgs)
-    | _      => return (getAllValuesForKey d k result, numArgs)
-
-partial def getUnify (d : DiscrTree α) (e : Expr) : MetaM (Array α) :=
-  withReducible do
-    let (k, args) ← getUnifyKeyArgs e (root := true)
-    match k with
-    | .star => d.root.foldlM (init := #[]) fun result k c => process k.arity #[] c result
-    | _ =>
-      let result := getStarResult d
-      match d.root.find? k with
-      | none   => return result
-      | 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
-            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)
-
-namespace Trie
-
-/--
-Monadically fold the keys and values stored in a `Trie`.
--/
-partial def foldM [Monad m] (initialKeys : Array Key)
-    (f : σ → Array Key → α → m σ) : (init : σ) → Trie α → m σ
-  | 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) =>
-      t.foldM (initialKeys.push k) f s
-
-/--
-Fold the keys and values stored in a `Trie`.
--/
-@[inline]
-def fold (initialKeys : Array Key) (f : σ → Array Key → α → σ) (init : σ) (t : Trie α) : σ :=
-  Id.run <| t.foldM initialKeys (init := init) fun s k a => return f s k a
-
-/--
-Monadically fold the values stored in a `Trie`.
--/
-partial def foldValuesM [Monad m] (f : σ → α → m σ) : (init : σ) → Trie α → m σ
-  | init, node vs children => do
-    let s ← vs.foldlM (init := init) f
-    children.foldlM (init := s) fun s (_, c) => c.foldValuesM (init := s) f
-
-/--
-Fold the values stored in a `Trie`.
--/
-@[inline]
-def foldValues (f : σ → α → σ) (init : σ) (t : Trie α) : σ :=
-  Id.run <| t.foldValuesM (init := init) (pure <| f · ·)
-
-/--
-The number of values stored in a `Trie`.
--/
-partial def size : Trie α → Nat
-  | Trie.node vs children =>
-    children.foldl (init := vs.size) fun n (_, c) => n + size c
-
-end Trie
-
-
-/--
-Monadically fold over the keys and values stored in a `DiscrTree`.
--/
-@[inline]
-def foldM [Monad m] (f : σ → Array Key → α → m σ) (init : σ)
-    (t : DiscrTree α) : m σ :=
-  t.root.foldlM (init := init) fun s k t => t.foldM #[k] (init := s) f
-
-/--
-Fold over the keys and values stored in a `DiscrTree`
--/
-@[inline]
-def fold (f : σ → Array Key → α → σ) (init : σ) (t : DiscrTree α) : σ :=
-  Id.run <| t.foldM (init := init) fun s keys a => return f s keys a
-
-/--
-Monadically fold over the values stored in a `DiscrTree`.
--/
-@[inline]
-def foldValuesM [Monad m] (f : σ → α → m σ) (init : σ) (t : DiscrTree α) :
-    m σ :=
-  t.root.foldlM (init := init) fun s _ t => t.foldValuesM (init := s) f
-
-/--
-Fold over the values stored in a `DiscrTree`.
--/
-@[inline]
-def foldValues (f : σ → α → σ) (init : σ) (t : DiscrTree α) : σ :=
-  Id.run <| t.foldValuesM (init := init) (pure <| f · ·)
-
-/--
-Check for the presence of a value satisfying a predicate.
--/
-@[inline]
-def containsValueP (t : DiscrTree α) (f : α → Bool) : Bool :=
-  t.foldValues (init := false) fun r a => r || f a
-
-/--
-Extract the values stored in a `DiscrTree`.
--/
-@[inline]
-def values (t : DiscrTree α) : Array α :=
-  t.foldValues (init := #[]) fun as a => as.push a
-
-/--
-Extract the keys and values stored in a `DiscrTree`.
--/
-@[inline]
-def toArray (t : DiscrTree α) : Array (Array Key × α) :=
-  t.fold (init := #[]) fun as keys a => as.push (keys, a)
-
-/--
-Get the number of values stored in a `DiscrTree`. O(n) in the size of the tree.
--/
-@[inline]
-def size (t : DiscrTree α) : Nat :=
-  t.root.foldl (init := 0) fun n _ t => n + t.size
-
-variable {m : Type → Type} [Monad m]
-
-/-- Apply a monadic function to the array of values at each node in a `DiscrTree`. -/
-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))
-
-/-- Apply a monadic function to the array of values at each node in a `DiscrTree`. -/
-def mapArraysM (d : DiscrTree α) (f : Array α → m (Array β)) : m (DiscrTree β) := do
-  pure { root := ← d.root.mapM (fun t => t.mapArraysM f) }
-
-/-- Apply a function to the array of values at each node in a `DiscrTree`. -/
-def mapArrays (d : DiscrTree α) (f : Array α → Array β) : DiscrTree β :=
-  Id.run <| d.mapArraysM fun A => pure (f A)
+public import Lean.Meta.DiscrTree.Basic
+public import Lean.Meta.DiscrTree.Util
+public import Lean.Meta.DiscrTree.Main

src/Lean/Meta/DiscrTree/Basic.lean

@@ -0,0 +1,118 @@
+/-
+Copyright (c) 2025 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.DiscrTree.Types
+public import Lean.ToExpr
+public import Lean.CoreM
+public section
+namespace Lean.Meta.DiscrTree
+
+instance : Inhabited (Trie α) where
+  default := .node #[] #[]
+
+instance : Inhabited (DiscrTree α) where
+  default := {}
+
+def empty : DiscrTree α := { root := {} }
+
+instance : ToExpr Key where
+  toTypeExpr := mkConst ``Key
+  toExpr k   := match k with
+   | .const n a => mkApp2 (mkConst ``Key.const) (toExpr n) (toExpr a)
+   | .fvar id a => mkApp2 (mkConst ``Key.fvar) (toExpr id) (toExpr a)
+   | .lit l => mkApp (mkConst ``Key.lit) (toExpr l)
+   | .star => mkConst ``Key.star
+   | .other => mkConst ``Key.other
+   | .arrow => mkConst ``Key.arrow
+   | .proj n i a => mkApp3 (mkConst ``Key.proj) (toExpr n) (toExpr i) (toExpr a)
+
+def Key.lt : Key → Key → Bool
+  | .lit v₁,        .lit v₂        => v₁ < v₂
+  | .fvar n₁ a₁,    .fvar n₂ a₂    => Name.quickLt n₁.name n₂.name || (n₁ == n₂ && a₁ < a₂)
+  | .const n₁ a₁,   .const n₂ a₂   => Name.quickLt n₁ n₂ || (n₁ == n₂ && a₁ < a₂)
+  | .proj s₁ i₁ a₁, .proj s₂ i₂ a₂ => Name.quickLt s₁ s₂ || (s₁ == s₂ && i₁ < i₂) || (s₁ == s₂ && i₁ == i₂ && a₁ < a₂)
+  | k₁,             k₂             => k₁.ctorIdx < k₂.ctorIdx
+
+instance : LT Key := ⟨fun a b => Key.lt a b⟩
+instance (a b : Key) : Decidable (a < b) := inferInstanceAs (Decidable (Key.lt a b))
+
+def Key.format : Key → Format
+  | .star            => "*"
+  | .other           => "◾"
+  | .lit (.natVal v) => Std.format v
+  | .lit (.strVal v) => repr v
+  | .const k _       => Std.format k
+  | .proj s i _      => Std.format s ++ "." ++ Std.format i
+  | .fvar k _        => Std.format k.name
+  | .arrow           => "∀"
+
+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))
+
+instance [ToFormat α] : ToFormat (Trie α) := ⟨Trie.format⟩
+
+partial def format [ToFormat α] (d : DiscrTree α) : Format :=
+  let (_, r) := d.root.foldl
+    (fun (p : Bool × Format) k c =>
+      (false, p.2 ++ (if p.1 then Format.nil else Format.line) ++ Format.paren (Std.format k ++ " => " ++ Std.format c)))
+    (true, Format.nil)
+  Format.group r
+
+instance [ToFormat α] : ToFormat (DiscrTree α) := ⟨format⟩
+
+/--
+Helper function for converting an entry (i.e., `Array Key`) to the discrimination tree into
+`MessageData` that is more user-friendly. We use this function to implement diagnostic information.
+-/
+partial def keysAsPattern (keys : Array Key) : CoreM MessageData := do
+  go (parenIfNonAtomic := false) |>.run' keys.toList
+where
+  next? : StateRefT (List Key) CoreM (Option Key) := do
+    let key :: keys ← get | return none
+    set keys
+    return some key
+
+  mkApp (f : MessageData) (args : Array MessageData) (parenIfNonAtomic : Bool) : CoreM MessageData := do
+    if args.isEmpty then
+      return f
+    else
+      let mut r := m!""
+      for arg in args do
+        r := r ++ Format.line ++ arg
+      r := f ++ .nest 2 r
+      if parenIfNonAtomic then
+        return .paren r
+      else
+        return .group r
+
+  go (parenIfNonAtomic := true) : StateRefT (List Key) CoreM MessageData := do
+    let some key ← next? | return .nil
+    match key with
+    | .const declName nargs =>
+      mkApp m!"{← mkConstWithLevelParams declName}" (← goN nargs) parenIfNonAtomic
+    | .fvar fvarId nargs =>
+      mkApp m!"{mkFVar fvarId}" (← goN nargs) parenIfNonAtomic
+    | .proj _ i nargs =>
+      mkApp m!"{← go}.{i+1}" (← goN nargs) parenIfNonAtomic
+    | .arrow =>
+      mkApp m!"∀ " (← goN 1) parenIfNonAtomic
+    | .star => return "_"
+    | .other => return "<other>"
+    | .lit (.natVal v) => return m!"{v}"
+    | .lit (.strVal v) => return m!"{v}"
+
+  goN (num : Nat) : StateRefT (List Key) CoreM (Array MessageData) := do
+    let mut r := #[]
+    for _ in *...num do
+      r := r.push (← go)
+    return r
+
+end Lean.Meta.DiscrTree

src/Lean/Meta/DiscrTree/Main.lean

@@ -0,0 +1,666 @@
+/-
+Copyright (c) 2019 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura, Jannis Limperg, Kim Morrison
+-/
+module
+prelude
+public import Lean.Meta.Basic
+public import Lean.Meta.DiscrTree.Basic
+import Lean.Meta.WHNF
+public section
+namespace Lean.Meta.DiscrTree
+/-!
+  (Imperfect) discrimination trees.
+  We use a hybrid representation.
+  - A `PersistentHashMap` for the root node which usually contains many children.
+  - A sorted array of key/node pairs for inner nodes.
+
+  The edges are labeled by keys:
+  - Constant names (and arity). Universe levels are ignored.
+  - Free variables (and arity). Thus, an entry in the discrimination tree
+    may reference hypotheses from the local context.
+  - Literals
+  - Star/Wildcard. We use them to represent metavariables and terms
+    we want to ignore. We ignore implicit arguments and proofs.
+  - Other. We use to represent other kinds of terms (e.g., nested lambda, forall, sort, etc).
+
+  We reduce terms using `TransparencyMode.reducible`. Thus, all reducible
+  definitions in an expression `e` are unfolded before we insert it into the
+  discrimination tree.
+
+  Recall that projections from classes are **NOT** reducible.
+  For example, the expressions `Add.add α (ringAdd ?α ?s) ?x ?x`
+  and `Add.add Nat Nat.hasAdd a b` generates paths with the following keys
+  respectively
+  ```
+  ⟨Add.add, 4⟩, α, *, *, *
+  ⟨Add.add, 4⟩, Nat, *, ⟨a,0⟩, ⟨b,0⟩
+  ```
+
+  That is, we don't reduce `Add.add Nat inst a b` into `Nat.add a b`.
+  We say the `Add.add` applications are the de-facto canonical forms in
+  the metaprogramming framework.
+  Moreover, it is the metaprogrammer's responsibility to re-pack applications such as
+  `Nat.add a b` into `Add.add Nat inst a b`.
+
+  Remark: we store the arity in the keys
+  1- To be able to implement the "skip" operation when retrieving "candidate"
+     unifiers.
+  2- Distinguish partial applications `f a`, `f a b`, and `f a b c`.
+-/
+
+def Key.arity : Key → Nat
+  | .const _ a  => a
+  | .fvar _ a   => a
+  /-
+  Remark: `.arrow` used to have arity 2, and was used to encode only **non**-dependent
+  arrows. However, this feature was a recurrent source of bugs. For example, a
+  theorem about a dependent arrow can be applied to a non-dependent one. The
+  reverse direction may also happen. See issue #2835. Therefore, `.arrow` was made
+  to have arity 0. But this throws away easy to use information, and makes it so
+  that ∀ and ∃ behave quite differently. So now `.arrow` at least indexes the
+  domain of the forall (whether dependent or non-dependent).
+  -/
+  | .arrow      => 1
+  | .proj _ _ a => 1 + a
+  | _           => 0
+
+/-- The discrimination tree ignores implicit arguments and proofs.
+   We use the following auxiliary id as a "mark". -/
+private def tmpMVarId : MVarId := { name := `_discr_tree_tmp }
+private def tmpStar := mkMVar tmpMVarId
+
+/--
+  Return true iff the argument should be treated as a "wildcard" by the discrimination tree.
+
+  - We ignore proofs because of proof irrelevance. It doesn't make sense to try to
+    index their structure.
+
+  - We ignore instance implicit arguments (e.g., `[Add α]`) because they are "morally" canonical.
+    Moreover, we may have many definitionally equal terms floating around.
+    Example: `Ring.hasAdd Int Int.isRing` and `Int.hasAdd`.
+
+  - We considered ignoring implicit arguments (e.g., `{α : Type}`) since users don't "see" them,
+    and may not even understand why some simplification rule is not firing.
+    However, in type class resolution, we have instance such as `Decidable (@Eq Nat x y)`,
+    where `Nat` is an implicit argument. Thus, we would add the path
+    ```
+    Decidable -> Eq -> * -> * -> * -> [Nat.decEq]
+    ```
+    to the discrimination tree IF we ignored the implicit `Nat` argument.
+    This would be BAD since **ALL** decidable equality instances would be in the same path.
+    So, we index implicit arguments if they are types.
+    This setting seems sensible for simplification theorems such as:
+    ```
+    forall (x y : Unit), (@Eq Unit x y) = true
+    ```
+    If we ignore the implicit argument `Unit`, the `DiscrTree` will say it is a candidate
+    simplification theorem for any equality in our goal.
+
+  Remark: if users have problems with the solution above, we may provide a `noIndexing` annotation,
+  and `ignoreArg` would return true for any term of the form `noIndexing t`.
+-/
+private def ignoreArg (a : Expr) (i : Nat) (infos : Array ParamInfo) : MetaM Bool := do
+  if h : i < infos.size then
+    let info := infos[i]
+    if info.isInstImplicit then
+      return true
+    else if info.isImplicit || info.isStrictImplicit then
+      return !(← isType a)
+    else
+      isProof a
+  else
+    isProof a
+
+private partial def pushArgsAux (infos : Array ParamInfo) : Nat → Expr → Array Expr → MetaM (Array Expr)
+  | i, .app f a, todo => do
+    if (← ignoreArg a i infos) then
+      pushArgsAux infos (i-1) f (todo.push tmpStar)
+    else
+      pushArgsAux infos (i-1) f (todo.push a)
+  | _, _, todo => return todo
+
+/--
+  Return true if `e` is one of the following
+  - A nat literal (numeral)
+  - `Nat.zero`
+  - `Nat.succ x` where `isNumeral x`
+  - `OfNat.ofNat _ x _` where `isNumeral x` -/
+private partial def isNumeral (e : Expr) : Bool :=
+  if e.isRawNatLit then true
+  else
+    let f := e.getAppFn
+    if !f.isConst then false
+    else
+      let fName := f.constName!
+      if fName == ``Nat.succ && e.getAppNumArgs == 1 then isNumeral e.appArg!
+      else if fName == ``OfNat.ofNat && e.getAppNumArgs == 3 then isNumeral (e.getArg! 1)
+      else if fName == ``Nat.zero && e.getAppNumArgs == 0 then true
+      else false
+
+private partial def toNatLit? (e : Expr) : Option Literal :=
+  if isNumeral e then
+    if let some n := loop e then
+      some (.natVal n)
+    else
+      none
+  else
+    none
+where
+  loop (e : Expr) : OptionT Id Nat := do
+    let f := e.getAppFn
+    match f with
+    | .lit (.natVal n) => return n
+    | .const fName .. =>
+      if fName == ``Nat.succ && e.getAppNumArgs == 1 then
+        let r ← loop e.appArg!
+        return r+1
+      else if fName == ``OfNat.ofNat && e.getAppNumArgs == 3 then
+        loop (e.getArg! 1)
+      else if fName == ``Nat.zero && e.getAppNumArgs == 0 then
+        return 0
+      else
+        failure
+    | _ => failure
+
+private def isNatType (e : Expr) : MetaM Bool :=
+  return (← whnf e).isConstOf ``Nat
+
+/--
+  Return true if `e` is one of the following
+  - `Nat.add _ k` where `isNumeral k`
+  - `Add.add Nat _ _ k` where `isNumeral k`
+  - `HAdd.hAdd _ Nat _ _ k` where `isNumeral k`
+  - `Nat.succ _`
+  This function assumes `e.isAppOf fName`
+-/
+private def isOffset (fName : Name) (e : Expr) : MetaM Bool := do
+  if fName == ``Nat.add && e.getAppNumArgs == 2 then
+    return isNumeral e.appArg!
+  else if fName == ``Add.add && e.getAppNumArgs == 4 then
+    if (← isNatType (e.getArg! 0)) then return isNumeral e.appArg! else return false
+  else if fName == ``HAdd.hAdd && e.getAppNumArgs == 6 then
+    if (← isNatType (e.getArg! 1)) then return isNumeral e.appArg! else return false
+  else
+    return fName == ``Nat.succ && e.getAppNumArgs == 1
+
+/--
+  TODO: add hook for users adding their own functions for controlling `shouldAddAsStar`
+  Different `DiscrTree` users may populate this set using, for example, attributes.
+
+  Remark: we currently tag "offset" terms as star to avoid having to add special
+  support for offset terms.
+  Example, suppose the discrimination tree contains the entry
+  `Nat.succ ?m |-> v`, and we are trying to retrieve the matches for `Expr.lit (Literal.natVal 1) _`.
+  In this scenario, we want to retrieve `Nat.succ ?m |-> v`
+-/
+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.
+-/
+partial def reduce (e : Expr) : MetaM Expr := do
+  let e ← whnfCore e
+  match (← unfoldDefinition? e) with
+  | some e => reduce e
+  | none => match e.etaExpandedStrict? with
+    | some e => reduce e
+    | none   => return e
+
+/--
+  Return `true` if `fn` is a "bad" key. That is, `pushArgs` would add `Key.other` or `Key.star`.
+  We use this function when processing "root terms, and will avoid unfolding terms.
+  Note that without this trick the pattern `List.map f ∘ List.map g` would be mapped into the key `Key.other`
+  since the function composition `∘` would be unfolded and we would get `fun x => List.map g (List.map f x)`
+-/
+private def isBadKey (fn : Expr) : Bool :=
+  match fn with
+  | .lit ..     => false
+  | .const ..   => false
+  | .fvar ..    => false
+  | .proj ..    => false
+  | .forallE .. => false
+  | _           => true
+
+/--
+Reduce `e` until we get an irreducible term (modulo current reducibility setting) or the resulting term
+is a bad key (see comment at `isBadKey`).
+We use this method instead of `reduce` for root terms at `pushArgs`. -/
+private partial def reduceUntilBadKey (e : Expr) : MetaM Expr := do
+  let e ← step e
+  match e.etaExpandedStrict? with
+  | some e => reduceUntilBadKey e
+  | none   => return e
+where
+  step (e : Expr) := do
+    let e ← whnfCore e
+    match (← unfoldDefinition? e) with
+    | some e' => if isBadKey e'.getAppFn then return e else step e'
+    | none    => return e
+
+/-- whnf for the discrimination tree module -/
+def reduceDT (e : Expr) (root : Bool) : MetaM Expr :=
+  if root then reduceUntilBadKey e else reduce e
+
+/- Remark: we use `shouldAddAsStar` only for nested terms, and `root == false` for nested terms -/
+
+/--
+Append `n` wildcards to `todo`
+-/
+private def pushWildcards (n : Nat) (todo : Array Expr) : Array Expr :=
+  match n with
+  | 0   => todo
+  | n+1 => pushWildcards n (todo.push tmpStar)
+
+/--
+When `noIndexAtArgs := true`, `pushArgs` assumes function application arguments have a `no_index` annotation.
+That is, `f a b` is indexed as it was `f (no_index a) (no_index b)`.
+This feature is used when indexing local proofs in the simplifier. This is useful in examples like the one described on issue #2670.
+In this issue, we have a local hypotheses `(h : ∀ p : α × β, f p p.2 = p.2)`, and users expect it to be applicable to
+`f (a, b) b = b`. This worked in Lean 3 since no indexing was used. We can retrieve Lean 3 behavior by writing
+`(h : ∀ p : α × β, f p (no_index p.2) = p.2)`, but this is very inconvenient when the hypotheses was not written by the user in first place.
+For example, it was introduced by another tactic. Thus, when populating the discrimination tree explicit arguments provided to `simp` (e.g., `simp [h]`),
+we use `noIndexAtArgs := true`. See comment: https://github.com/leanprover/lean4/issues/2670#issuecomment-1758889365
+-/
+private def pushArgs (root : Bool) (todo : Array Expr) (e : Expr) (noIndexAtArgs : Bool) : MetaM (Key × Array Expr) := do
+  if hasNoindexAnnotation e then
+    return (.star, todo)
+  else
+    let e ← reduceDT e root
+    let fn := e.getAppFn
+    let push (k : Key) (nargs : Nat) (todo : Array Expr): MetaM (Key × Array Expr) := do
+      let info ← getFunInfoNArgs fn nargs
+      let todo ← if noIndexAtArgs then
+        pure <| pushWildcards nargs todo
+      else
+        pushArgsAux info.paramInfo (nargs-1) e todo
+      return (k, todo)
+    match fn with
+    | .lit v     =>
+      return (.lit v, todo)
+    | .const c _ =>
+      unless root do
+        if let some v := toNatLit? e then
+          return (.lit v, todo)
+        if (← shouldAddAsStar c e) then
+          return (.star, todo)
+      let nargs := e.getAppNumArgs
+      push (.const c nargs) nargs todo
+    | .proj s i a =>
+      /-
+      If `s` is a class, then `a` is an instance. Thus, we annotate `a` with `no_index` since we do not
+      index instances. This should only happen if users mark a class projection function as `[reducible]`.
+
+      TODO: add better support for projections that are functions
+      -/
+      let a := if isClass (← getEnv) s then mkNoindexAnnotation a else a
+      let nargs := e.getAppNumArgs
+      push (.proj s i nargs) nargs (todo.push a)
+    | .fvar fvarId   =>
+      let nargs := e.getAppNumArgs
+      push (.fvar fvarId nargs) nargs todo
+    | .mvar mvarId   =>
+      if mvarId == tmpMVarId then
+        -- We use `tmp to mark implicit arguments and proofs
+        return (.star, todo)
+      else if (← mvarId.isReadOnlyOrSyntheticOpaque) then
+        return (.other, todo)
+      else
+        return (.star, todo)
+    | .forallE _n d _ _ =>
+      return (.arrow, todo.push d)
+    | _ => return (.other, todo)
+
+@[inherit_doc pushArgs]
+partial def mkPathAux (root : Bool) (todo : Array Expr) (keys : Array Key) (noIndexAtArgs : Bool) : MetaM (Array Key) := do
+  if todo.isEmpty then
+    return keys
+  else
+    let e    := todo.back!
+    let todo := todo.pop
+    let (k, todo) ← pushArgs root todo e noIndexAtArgs
+    mkPathAux false todo (keys.push k) noIndexAtArgs
+
+private def initCapacity := 8
+
+@[inherit_doc pushArgs]
+def mkPath (e : Expr) (noIndexAtArgs := false) : MetaM (Array Key) := do
+  withReducible do
+    let todo : Array Expr := .mkEmpty initCapacity
+    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
+
+/--
+Inserts a value into a discrimination tree,
+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
+  else
+    d.insertCore keys v
+
+private def getKeyArgs (e : Expr) (isMatch root : Bool) : MetaM (Key × Array Expr) := do
+  let e ← reduceDT e root
+  unless root do
+    -- See pushArgs
+    if let some v := toNatLit? e then
+      return (.lit v, #[])
+  match e.getAppFn with
+  | .lit v         => return (.lit v, #[])
+  | .const c _     =>
+    if (← getConfig).isDefEqStuckEx && e.hasExprMVar then
+      if (← isReducible c) then
+        /- `e` is a term `c ...` s.t. `c` is reducible and `e` has metavariables, but it was not unfolded.
+           This can happen if the metavariables in `e` are "blocking" smart unfolding.
+           If `isDefEqStuckEx` is enabled, then we must throw the `isDefEqStuck` exception to postpone TC resolution.
+           Here is an example. Suppose we have
+           ```
+            inductive Ty where
+              | bool | fn (a ty : Ty)
+
+
+            @[reducible] def Ty.interp : Ty → Type
+              | bool   => Bool
+              | fn a b => a.interp → b.interp
+           ```
+           and we are trying to synthesize `BEq (Ty.interp ?m)`
+        -/
+        Meta.throwIsDefEqStuck
+      else if let some matcherInfo := isMatcherAppCore? (← getEnv) e then
+        -- A matcher application is stuck is one of the discriminants has a metavariable
+        let args := e.getAppArgs
+        for arg in args[matcherInfo.getFirstDiscrPos...(matcherInfo.getFirstDiscrPos + matcherInfo.numDiscrs)] do
+          if arg.hasExprMVar then
+            Meta.throwIsDefEqStuck
+      else if (← isRec c) then
+        /- Similar to the previous case, but for `match` and recursor applications. It may be stuck (i.e., did not reduce)
+           because of metavariables. -/
+        Meta.throwIsDefEqStuck
+    let nargs := e.getAppNumArgs
+    return (.const c nargs, e.getAppRevArgs)
+  | .fvar fvarId   =>
+    let nargs := e.getAppNumArgs
+    return (.fvar fvarId nargs, e.getAppRevArgs)
+  | .mvar mvarId   =>
+    if isMatch then
+      return (.other, #[])
+    else do
+      let cfg ← getConfig
+      if cfg.isDefEqStuckEx then
+        /-
+          When the configuration flag `isDefEqStuckEx` is set to true,
+          we want `isDefEq` to throw an exception whenever it tries to assign
+          a read-only metavariable.
+          This feature is useful for type class resolution where
+          we may want to notify the caller that the TC problem may be solvable
+          later after it assigns `?m`.
+          The method `DiscrTree.getUnify e` returns candidates `c` that may "unify" with `e`.
+          That is, `isDefEq c e` may return true. Now, consider `DiscrTree.getUnify d (Add ?m)`
+          where `?m` is a read-only metavariable, and the discrimination tree contains the keys
+          `HadAdd Nat` and `Add Int`. If `isDefEqStuckEx` is set to true, we must treat `?m` as
+          a regular metavariable here, otherwise we return the empty set of candidates.
+          This is incorrect because it is equivalent to saying that there is no solution even if
+          the caller assigns `?m` and try again. -/
+        return (.star, #[])
+      else if (← mvarId.isReadOnlyOrSyntheticOpaque) then
+        return (.other, #[])
+      else
+        return (.star, #[])
+  | .proj s i a .. =>
+    let nargs := e.getAppNumArgs
+    return (.proj s i nargs, #[a] ++ e.getAppRevArgs)
+  | .forallE _ d _ _ => return (.arrow, #[d])
+  | _ => return (.other, #[])
+
+private abbrev getMatchKeyArgs (e : Expr) (root : Bool) : MetaM (Key × Array Expr) :=
+  getKeyArgs e (isMatch := true) (root := root)
+
+private abbrev getUnifyKeyArgs (e : Expr) (root : Bool) : MetaM (Key × Array Expr) :=
+  getKeyArgs e (isMatch := false) (root := root)
+
+private def getStarResult (d : DiscrTree α) : Array α :=
+  let result : Array α := .mkEmpty initCapacity
+  match d.root.find? .star with
+  | none                  => result
+  | 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
+  | .node vs cs =>
+    if todo.isEmpty then
+      return result ++ vs
+    else if cs.isEmpty then
+      return result
+    else
+      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`. -/
+      let visitStar (result : Array α) : MetaM (Array α) :=
+        if first.1 == .star then
+          getMatchLoop 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 => 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
+  | none   => return result
+  | some c => getMatchLoop args c result
+
+private def getMatchCore (d : DiscrTree α) (e : Expr) : MetaM (Key × Array α) :=
+  withReducible do
+    let result := getStarResult d
+    let (k, args) ← getMatchKeyArgs e (root := true)
+    match k with
+    | .star  => return (k, result)
+    | _      => return (k, (← getMatchRoot d k args result))
+
+/--
+  Find values that match `e` in `d`.
+-/
+def getMatch (d : DiscrTree α) (e : Expr) : MetaM (Array α) :=
+  return (← getMatchCore d e).2
+
+/--
+  Similar to `getMatch`, but returns solutions that are prefixes of `e`.
+  We store the number of ignored arguments in the result.-/
+partial def getMatchWithExtra (d : DiscrTree α) (e : Expr) : MetaM (Array (α × Nat)) := do
+  let (k, result) ← getMatchCore d e
+  let result := result.map (·, 0)
+  if !e.isApp then
+    return result
+  else if !(← mayMatchPrefix k) then
+    return result
+  else
+    go e.appFn! 1 result
+where
+  mayMatchPrefix (k : Key) : MetaM Bool :=
+    let cont (k : Key) : MetaM Bool :=
+      if d.root.find? k |>.isSome then
+        return true
+      else
+        mayMatchPrefix k
+    match k with
+    | .const f (n+1)  => cont (.const f n)
+    | .fvar f (n+1)   => cont (.fvar f n)
+    | .proj s i (n+1) => cont (.proj s i n)
+    | _               => return false
+
+  go (e : Expr) (numExtra : Nat) (result : Array (α × Nat)) : MetaM (Array (α × Nat)) := do
+    let result := result ++ (← getMatchCore d e).2.map (., numExtra)
+    if e.isApp then
+      go e.appFn! (numExtra + 1) result
+    else
+      return result
+
+/--
+Return the root symbol for `e`, and the number of arguments after `reduceDT`.
+-/
+def getMatchKeyRootFor (e : Expr) : MetaM (Key × Nat) := do
+  let e ← reduceDT e (root := true)
+  let numArgs := e.getAppNumArgs
+  let key := match e.getAppFn with
+    | .lit v         => .lit v
+    | .fvar fvarId   => .fvar fvarId numArgs
+    | .mvar _        => .other
+    | .proj s i _ .. => .proj s i numArgs
+    | .forallE ..    => .arrow
+    | .const c _     =>
+      -- This method is used by the simplifier only, we do **not** support
+      -- (← getConfig).isDefEqStuckEx
+      .const c numArgs
+    | _ => .other
+  return (key, numArgs)
+
+/--
+Get all results under key `k`.
+-/
+private partial def getAllValuesForKey (d : DiscrTree α) (k : Key) (result : Array α) : Array α :=
+  match d.root.find? k with
+  | none      => result
+  | some trie => go trie result
+where
+  go (trie : Trie α) (result : Array α) : Array α := Id.run do
+    match trie with
+    | .node vs cs =>
+      let mut result := result ++ vs
+      for (_, trie) in cs do
+        result := go trie result
+      return result
+
+/--
+A liberal version of `getMatch` which only takes the root symbol of `e` into account.
+We use this method to simulate Lean 3's indexing.
+
+The natural number in the result is the number of arguments in `e` after `reduceDT`.
+-/
+def getMatchLiberal (d : DiscrTree α) (e : Expr) : MetaM (Array α × Nat) := do
+  withReducible do
+    let result := getStarResult d
+    let (k, numArgs) ← getMatchKeyRootFor e
+    match k with
+    | .star  => return (result, numArgs)
+    | _      => return (getAllValuesForKey d k result, numArgs)
+
+partial def getUnify (d : DiscrTree α) (e : Expr) : MetaM (Array α) :=
+  withReducible do
+    let (k, args) ← getUnifyKeyArgs e (root := true)
+    match k with
+    | .star => d.root.foldlM (init := #[]) fun result k c => process k.arity #[] c result
+    | _ =>
+      let result := getStarResult d
+      match d.root.find? k with
+      | none   => return result
+      | 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
+            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)
+
+end Lean.Meta.DiscrTree

src/Lean/Meta/DiscrTree/Types.lean

@@ -4,16 +4,11 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.ToExpr
-
 public section
-
 namespace Lean.Meta
-
 /-! See file `DiscrTree.lean` for the actual implementation and documentation. -/
-
 namespace DiscrTree
 /--
 Discrimination tree key. See `DiscrTree`
@@ -39,17 +34,6 @@ protected def Key.hash : Key → UInt64
 
 instance : Hashable Key := ⟨Key.hash⟩
 
-instance : ToExpr Key where
-  toTypeExpr := mkConst ``Key
-  toExpr k   := match k with
-   | .const n a => mkApp2 (mkConst ``Key.const) (toExpr n) (toExpr a)
-   | .fvar id a => mkApp2 (mkConst ``Key.fvar) (toExpr id) (toExpr a)
-   | .lit l => mkApp (mkConst ``Key.lit) (toExpr l)
-   | .star => mkConst ``Key.star
-   | .other => mkConst ``Key.other
-   | .arrow => mkConst ``Key.arrow
-   | .proj n i a => mkApp3 (mkConst ``Key.proj) (toExpr n) (toExpr i) (toExpr a)
-
 /--
 Discrimination tree trie. See `DiscrTree`.
 -/

src/Lean/Meta/DiscrTree/Util.lean

@@ -0,0 +1,129 @@
+/-
+Copyright (c) 2024 Lean FRO, LLC. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Kim Morrison
+-/
+module
+prelude
+public import Lean.Meta.DiscrTree.Basic
+public section
+namespace Lean.Meta.DiscrTree
+namespace Trie
+/--
+Monadically fold the keys and values stored in a `Trie`.
+-/
+partial def foldM [Monad m] (initialKeys : Array Key)
+    (f : σ → Array Key → α → m σ) : (init : σ) → Trie α → m σ
+  | 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) =>
+      t.foldM (initialKeys.push k) f s
+
+/--
+Fold the keys and values stored in a `Trie`.
+-/
+@[inline]
+def fold (initialKeys : Array Key) (f : σ → Array Key → α → σ) (init : σ) (t : Trie α) : σ :=
+  Id.run <| t.foldM initialKeys (init := init) fun s k a => return f s k a
+
+/--
+Monadically fold the values stored in a `Trie`.
+-/
+partial def foldValuesM [Monad m] (f : σ → α → m σ) : (init : σ) → Trie α → m σ
+  | init, node vs children => do
+    let s ← vs.foldlM (init := init) f
+    children.foldlM (init := s) fun s (_, c) => c.foldValuesM (init := s) f
+
+/--
+Fold the values stored in a `Trie`.
+-/
+@[inline]
+def foldValues (f : σ → α → σ) (init : σ) (t : Trie α) : σ :=
+  Id.run <| t.foldValuesM (init := init) (pure <| f · ·)
+
+/--
+The number of values stored in a `Trie`.
+-/
+partial def size : Trie α → Nat
+  | Trie.node vs children =>
+    children.foldl (init := vs.size) fun n (_, c) => n + size c
+
+end Trie
+
+
+/--
+Monadically fold over the keys and values stored in a `DiscrTree`.
+-/
+@[inline]
+def foldM [Monad m] (f : σ → Array Key → α → m σ) (init : σ)
+    (t : DiscrTree α) : m σ :=
+  t.root.foldlM (init := init) fun s k t => t.foldM #[k] (init := s) f
+
+/--
+Fold over the keys and values stored in a `DiscrTree`
+-/
+@[inline]
+def fold (f : σ → Array Key → α → σ) (init : σ) (t : DiscrTree α) : σ :=
+  Id.run <| t.foldM (init := init) fun s keys a => return f s keys a
+
+/--
+Monadically fold over the values stored in a `DiscrTree`.
+-/
+@[inline]
+def foldValuesM [Monad m] (f : σ → α → m σ) (init : σ) (t : DiscrTree α) :
+    m σ :=
+  t.root.foldlM (init := init) fun s _ t => t.foldValuesM (init := s) f
+
+/--
+Fold over the values stored in a `DiscrTree`.
+-/
+@[inline]
+def foldValues (f : σ → α → σ) (init : σ) (t : DiscrTree α) : σ :=
+  Id.run <| t.foldValuesM (init := init) (pure <| f · ·)
+
+/--
+Check for the presence of a value satisfying a predicate.
+-/
+@[inline]
+def containsValueP (t : DiscrTree α) (f : α → Bool) : Bool :=
+  t.foldValues (init := false) fun r a => r || f a
+
+/--
+Extract the values stored in a `DiscrTree`.
+-/
+@[inline]
+def values (t : DiscrTree α) : Array α :=
+  t.foldValues (init := #[]) fun as a => as.push a
+
+/--
+Extract the keys and values stored in a `DiscrTree`.
+-/
+@[inline]
+def toArray (t : DiscrTree α) : Array (Array Key × α) :=
+  t.fold (init := #[]) fun as keys a => as.push (keys, a)
+
+/--
+Get the number of values stored in a `DiscrTree`. O(n) in the size of the tree.
+-/
+@[inline]
+def size (t : DiscrTree α) : Nat :=
+  t.root.foldl (init := 0) fun n _ t => n + t.size
+
+variable {m : Type → Type} [Monad m]
+
+/-- Apply a monadic function to the array of values at each node in a `DiscrTree`. -/
+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))
+
+/-- Apply a monadic function to the array of values at each node in a `DiscrTree`. -/
+def mapArraysM (d : DiscrTree α) (f : Array α → m (Array β)) : m (DiscrTree β) := do
+  pure { root := ← d.root.mapM (fun t => t.mapArraysM f) }
+
+/-- Apply a function to the array of values at each node in a `DiscrTree`. -/
+def mapArrays (d : DiscrTree α) (f : Array α → Array β) : DiscrTree β :=
+  Id.run <| d.mapArraysM fun A => pure (f A)
+
+end Lean.Meta.DiscrTree

src/Lean/Meta/ExprDefEq.lean

@@ -4,13 +4,11 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Meta.UnificationHint
 public import Lean.Util.OccursCheck
-
+import Lean.Meta.WHNF
 public section
-
 namespace Lean.Meta
 
 register_builtin_option backward.isDefEq.lazyProjDelta : Bool := {

src/Lean/Meta/Instances.lean

@@ -4,14 +4,12 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Init.Data.Range.Polymorphic.Stream
-public import Lean.Meta.DiscrTree
+public import Lean.Meta.DiscrTree.Main
 public import Lean.Meta.CollectMVars
-
+import Lean.Meta.WHNF
 public section
-
 namespace Lean.Meta
 
 register_builtin_option synthInstance.checkSynthOrder : Bool := {

src/Lean/Meta/Match/NamedPatterns.lean

@@ -4,15 +4,14 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Meta.Basic
 import Lean.Meta.AppBuilder
-
+import Lean.Meta.Transform
+import Lean.Meta.WHNF
 /-!
 Helper functions around named patterns
 -/
-
 namespace Lean.Meta.Match
 
 public def mkNamedPattern (x h p : Expr) : MetaM Expr :=
@@ -36,3 +35,5 @@ public def unfoldNamedPattern (e : Expr) : MetaM Expr := do
         return TransformStep.visit eNew
     return .continue
   Meta.transform e (pre := visit)
+
+end Lean.Meta.Match

src/Lean/Meta/MethodSpecs.lean

@@ -3,12 +3,11 @@ Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Joachim Breitner
 -/
-
 module
 prelude
 public import Lean.Meta.Tactic.Simp.SimpTheorems
 import Lean.Meta.Tactic.Simp.Main
-
+import Lean.Structure
 namespace Lean
 
 open Meta

src/Lean/Meta/SizeOf.lean

@@ -8,6 +8,7 @@ prelude
 public import Lean.AddDecl
 public import Lean.Meta.AppBuilder
 public import Lean.DefEqAttrib
+import Lean.Meta.WHNF
 public section
 namespace Lean.Meta
 

src/Lean/Meta/Structure.lean

@@ -4,19 +4,17 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura, Kyle Miller
 -/
 module
-
 prelude
 public import Lean.AddDecl
 public import Lean.Meta.AppBuilder
-
+import Lean.Structure
+import Lean.Meta.Transform
 public section
-
+namespace Lean.Meta
 /-!
 # Structure methods that require `MetaM` infrastructure
 -/
 
-namespace Lean.Meta
-
 /--
 If `struct` is an application of the form `S ..` with `S` a constant for a structure,
 returns the name of the structure, otherwise throws an error.

src/Lean/Meta/Tactic/CasesOnStuckLHS.lean

@@ -4,11 +4,10 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Meta.Basic
 import Lean.Meta.Tactic.SplitIf
-
+import Lean.ProjFns
 namespace Lean.Meta
 
 /-!

src/Lean/Meta/Tactic/Delta.lean

@@ -4,12 +4,10 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Meta.Tactic.Replace
-
+import Lean.Meta.Transform
 public section
-
 namespace Lean.Meta
 
 def delta? (e : Expr) (p : Name → Bool := fun _ => true) : CoreM (Option Expr) :=

src/Lean/Meta/Tactic/Ext.lean

@@ -4,13 +4,10 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Gabriel Ebner, Mario Carneiro
 -/
 module
-
 prelude
 public import Init.Data.Array.InsertionSort
 public import Lean.Meta.DiscrTree
-
 public section
-
 namespace Lean.Meta.Ext
 
 /-!

src/Lean/Meta/Tactic/Grind/Arith/Cutsat/ToIntInfo.lean

@@ -7,8 +7,8 @@ module
 prelude
 public import Lean.Meta.Tactic.Grind.Arith.Util
 import Init.Grind.ToInt
+import Lean.Meta.LitValues
 public section
-
 namespace Lean.Meta.Grind.Arith.Cutsat
 open Lean Grind
 

src/Lean/Meta/Tactic/Grind/Injection.lean

@@ -8,9 +8,8 @@ prelude
 public import Lean.Meta.Basic
 import Lean.Meta.Tactic.Clear
 import Lean.Meta.AppBuilder
-
+import Lean.Meta.CtorRecognizer
 public section
-
 namespace Lean.Meta.Grind
 /--
 The `grind` tactic includes an auxiliary `injection?` tactic that is not intended for direct use by users.

src/Lean/Meta/Tactic/Grind/Util.lean

@@ -7,6 +7,8 @@ module
 prelude
 public import Lean.Meta.Tactic.Simp.Simproc
 import Init.Simproc
+import Lean.ProjFns
+import Lean.Meta.WHNF
 import Lean.Meta.AbstractNestedProofs
 import Lean.Meta.Tactic.Clear
 public section

src/Lean/Meta/Tactic/Induction.lean

@@ -4,16 +4,14 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Meta.RecursorInfo
 public import Lean.Meta.SynthInstance
 public import Lean.Meta.Tactic.Revert
 public import Lean.Meta.Tactic.Intro
 public import Lean.Meta.Tactic.FVarSubst
-
+import Lean.Meta.WHNF
 public section
-
 namespace Lean.Meta
 
 private partial def getTargetArity : Expr → Nat

src/Lean/Meta/Tactic/Simp/Arith/Int/Basic.lean

@@ -4,7 +4,6 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Init.Data.Int.Linear
 public import Lean.Util.SortExprs
@@ -12,9 +11,8 @@ public import Lean.Meta.IntInstTesters
 public import Lean.Meta.AppBuilder
 public import Lean.Meta.KExprMap
 public import Lean.Data.RArray
-
+import Lean.Meta.LitValues
 public section
-
 namespace Int.Linear
 
 /-- Converts the linear polynomial into the "simplified" expression -/

src/Lean/Meta/Tactic/Simp/Arith/Nat/Basic.lean

@@ -10,6 +10,7 @@ public import Lean.Meta.KExprMap
 import Lean.Data.RArray
 import Lean.Meta.AppBuilder
 import Lean.Meta.NatInstTesters
+import Lean.Meta.Offset
 public section
 namespace Nat.Linear
 

src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/Core.lean

@@ -4,13 +4,11 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Init.Simproc
 public import Lean.Meta.Tactic.Simp.Simproc
-
+import Lean.Meta.CtorRecognizer
 public section
-
 open Lean Meta Simp
 
 builtin_simproc ↓ [simp, seval] reduceIte (ite _ _ _) := fun e => do

src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/CtorIdx.lean

@@ -3,16 +3,13 @@ Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Joachim Breitner
 -/
-
 module
-
 prelude
 public import Lean.Meta.Tactic.Simp.Simproc
 import Init.Simproc
 import Lean.Meta.Constructions.CtorIdx
-
+import Lean.Meta.CtorRecognizer
 open Lean Meta Simp
-
 public section
 
 /--

src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/Int.lean

@@ -4,12 +4,10 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Meta.Tactic.Simp.BuiltinSimprocs.Nat
-
+import Lean.Util.SafeExponentiation
 public section
-
 namespace Int
 open Lean Meta Simp
 

src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/Nat.lean

@@ -4,13 +4,13 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Init.Simproc
 public import Lean.Meta.Tactic.Simp.BuiltinSimprocs.Util
-
+public import Lean.Meta.LitValues
+public import Lean.Meta.Offset
+import Lean.Util.SafeExponentiation
 public section
-
 namespace Nat
 open Lean Meta Simp
 

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

@@ -4,7 +4,6 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Meta.ACLt
 public import Lean.Meta.Match.MatchEqsExt
@@ -12,11 +11,10 @@ public import Lean.Meta.Tactic.UnifyEq
 public import Lean.Meta.Tactic.Simp.Arith
 public import Lean.Meta.Tactic.Simp.Attr
 public import Lean.Meta.BinderNameHint
-
+import Lean.Meta.FunInfo
+import Lean.Meta.WHNF
 public section
-
 namespace Lean.Meta.Simp
-
 /--
 Helper type for implementing `discharge?'`
 -/

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

@@ -4,15 +4,15 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-public import Lean.Meta.DiscrTree
+public import Lean.Meta.DiscrTree.Main
 public import Lean.Meta.Tactic.AuxLemma
 public import Lean.DocString
+import Lean.ExtraModUses
+import Lean.ProjFns
 import Lean.Meta.AppBuilder
 import Lean.Meta.Eqns
-import Lean.ExtraModUses
-
+import Lean.Meta.WHNF
 public section
 
 /-!

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

@@ -4,7 +4,6 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Meta.AppBuilder
 public import Lean.Meta.CongrTheorems
@@ -12,9 +11,8 @@ public import Lean.Meta.Eqns
 public import Lean.Meta.Tactic.Simp.SimpTheorems
 public import Lean.Meta.Tactic.Simp.SimpCongrTheorems
 import Lean.Meta.Tactic.Replace
-
+import Lean.Meta.FunInfo
 public section
-
 namespace Lean.Meta
 namespace Simp
 

src/Lean/Meta/Tactic/SplitIf.lean

@@ -4,14 +4,12 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Meta.Tactic.Cases
 public import Lean.Meta.Tactic.Simp.Rewrite
 import Lean.Meta.Tactic.Simp.Main
-
+import Lean.Meta.FunInfo
 public section
-
 namespace Lean.Meta
 
 inductive SplitKind where

src/Lean/Meta/Tactic/Symm.lean

@@ -4,13 +4,11 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Siddhartha Gadgil, Mario Carneiro, Kim Morrison
 -/
 module
-
 prelude
 public import Lean.Meta.Reduce
 public import Lean.Meta.Tactic.Assert
-public import Lean.Meta.DiscrTree
+public import Lean.Meta.DiscrTree.Main
 import Lean.Meta.AppBuilder
-
 public section
 
 /-!

src/Lean/Meta/Tactic/Unfold.lean

@@ -4,13 +4,11 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Meta.Tactic.Delta
 public import Lean.Meta.Tactic.Simp.Main
-
+import Lean.Meta.WHNF
 public section
-
 namespace Lean.Meta
 
 private def getSimpUnfoldContext : MetaM Simp.Context := do