refactor: replace fullExt with starLemmasExt and add eliminator support

This PR makes two main improvements to library search:

1. **Replace duplicate tree with array for star-indexed lemmas**
   - Instead of maintaining `fullExt` (a complete duplicate discrimination tree),
     use `starLemmasExt` (a simple array for lemmas with `[*]` or `[Eq,*,*,*]` keys)
   - Star-indexed lemmas are captured during tree initialization via `extractKeys`
   - Two-pass search: first excludes star lemmas, fallback includes them
   - Controlled by `-star`/`+star` flags

2. **Support eliminator-style theorems like `iteInduction`**
   - Add `getFirstArgEntry` to create secondary index entries for fvar-headed theorems
   - For `motive (ite c t e)`, creates entry keyed by `.const ite 4`
   - Modify `getMatchCore` to check first argument's const key for fvar-headed goals
   - This enables finding eliminators even with `-star`

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

Author: kim-em

src/Lean/Meta/LazyDiscrTree.lean

@@ -685,11 +685,21 @@ def getMatchCore (root : Std.HashMap Key TrieIndex) (e : Expr) :
     | .star  =>
       pure #[]
     /- When goal has fvar head like `p (ite c t e)`, follow the star edge with the fvar's arguments.
-       This finds "eliminator-style" theorems indexed by argument structure. -/
-    | .fvar _ _ =>
-      match root[Key.star]? with
-      | some c => pure #[{ todo := args, score := 0, c }]
-      | none => pure #[]
+       This finds "eliminator-style" theorems indexed by argument structure.
+       We also check the first argument's const key directly, which enables finding
+       eliminator-style theorems even with -star (when star entries are dropped). -/
+    | .fvar _ _ => do
+      let mut cases := #[]
+      -- Follow star edge if available (for generic fvar-headed lemmas)
+      if let some c := root[Key.star]? then
+        cases := cases.push { todo := args, score := 0, c }
+      -- Also check first argument's const key (for eliminator-style theorems)
+      if !args.isEmpty then
+        let firstArg := args[0]!
+        let (argKey, argArgs) ← MatchClone.getMatchKeyArgs firstArg (root := false)
+        if let some c := root[argKey]? then
+          cases := cases.push { todo := argArgs, score := 1, c }
+      pure cases
     /- See note about "dep-arrow vs arrow" at `getMatchLoop` -/
     | .arrow =>
       pure (#[] |> pushRootCase root .other #[]
@@ -939,6 +949,57 @@ def createLocalPreDiscrTree
 def dropKeys (t : LazyDiscrTree α) (keys : List (List LazyDiscrTree.Key)) : MetaM (LazyDiscrTree α) := do
   keys.foldlM (init := t) (·.dropKey ·)
 
+/-- Collect all values from a subtree recursively and clear them. -/
+partial def collectSubtreeAux (next : TrieIndex) : MatchM α (Array α) :=
+  if next = 0 then
+    pure #[]
+  else do
+    let (values, star, children) ← evalNode next
+    -- Collect from star subtrie
+    let starVals ← collectSubtreeAux star
+    -- Collect from all children
+    let mut childVals : Array α := #[]
+    for (_, childIdx) in children do
+      childVals := childVals ++ (← collectSubtreeAux childIdx)
+    -- Clear this node (keep structure but remove values)
+    modify (·.set! next {values := #[], star, children})
+    return values ++ starVals ++ childVals
+
+/-- Navigate to a key path and return all values in that subtree, then drop them. -/
+def extractKeyAux (next : TrieIndex) (rest : List Key) :
+    MatchM α (Array α) :=
+  if next = 0 then
+    pure #[]
+  else do
+    let (_, star, children) ← evalNode next
+    match rest with
+    | [] =>
+      -- At the target node: collect ALL values from entire subtree
+      collectSubtreeAux next
+    | k :: r => do
+      let next := if k == .star then star else children.getD k 0
+      extractKeyAux next r
+
+/-- Extract and drop entries at a specific key, returning the dropped entries. -/
+def extractKey (t : LazyDiscrTree α) (path : List LazyDiscrTree.Key) :
+    MetaM (Array α × LazyDiscrTree α) :=
+  match path with
+  | [] => pure (#[], t)
+  | rootKey :: rest => do
+    let idx := t.roots.getD rootKey 0
+    runMatch t (extractKeyAux idx rest)
+
+/-- Extract entries at the given keys and also drop them from the tree. -/
+def extractKeys (t : LazyDiscrTree α) (keys : List (List LazyDiscrTree.Key)) :
+    MetaM (Array α × LazyDiscrTree α) := do
+  let mut allExtracted : Array α := #[]
+  let mut tree := t
+  for path in keys do
+    let (extracted, newTree) ← extractKey tree path
+    allExtracted := allExtracted ++ extracted
+    tree := newTree
+  return (allExtracted, tree)
+
 def logImportFailure [Monad m] [MonadLog m] [AddMessageContext m] [MonadOptions m] (f : ImportFailure) : m Unit :=
   logError m!"Processing failure with {f.const} in {f.module}:\n  {f.exception.toMessageData}"
 
@@ -990,6 +1051,7 @@ def findImportMatches
       (addEntry : Name → ConstantInfo → MetaM (Array (InitEntry α)))
       (droppedKeys : List (List LazyDiscrTree.Key) := [])
       (constantsPerTask : Nat := 1000)
+      (droppedEntriesRef : Option (IO.Ref (Option (Array α))) := none)
       (ty : Expr) : MetaM (MatchResult α) := do
   let cctx ← (read : CoreM Core.Context)
   let ngen ← getNGen
@@ -1001,7 +1063,13 @@ def findImportMatches
     profileitM Exception  "lazy discriminator import initialization" (←getOptions) $ do
       let t ← createImportedDiscrTree (createTreeCtx cctx) cNGen (←getEnv) addEntry
                 (constantsPerTask := constantsPerTask)
-      dropKeys t droppedKeys
+      -- If a reference is provided, extract and store dropped entries
+      if let some droppedRef := droppedEntriesRef then
+        let (extracted, t) ← extractKeys t droppedKeys
+        droppedRef.set (some extracted)
+        pure t
+      else
+        dropKeys t droppedKeys
   let (importCandidates, importTree) ← importTree.getMatch ty
   ref.set (some importTree)
   pure importCandidates
@@ -1075,10 +1143,11 @@ def findMatchesExt
     (addEntry : Name → ConstantInfo → MetaM (Array (InitEntry α)))
     (droppedKeys : List (List LazyDiscrTree.Key) := [])
     (constantsPerTask : Nat := 1000)
+    (droppedEntriesRef : Option (IO.Ref (Option (Array α))) := none)
     (adjustResult : Nat → α → β)
     (ty : Expr) : MetaM (Array β) := do
   let moduleMatches ← findModuleMatches moduleTreeRef ty
-  let importMatches ← findImportMatches ext addEntry droppedKeys constantsPerTask ty
+  let importMatches ← findImportMatches ext addEntry droppedKeys constantsPerTask droppedEntriesRef ty
   return Array.mkEmpty (moduleMatches.size + importMatches.size)
           |> moduleMatches.appendResultsAux (f := adjustResult)
           |> importMatches.appendResultsAux (f := adjustResult)
@@ -1091,13 +1160,15 @@ def findMatchesExt
 * `addEntry` is the function for creating discriminator tree entries from constants.
 * `droppedKeys` contains keys we do not want to consider when searching for matches.
   It is used for dropping very general keys.
+* `droppedEntriesRef` optionally stores entries dropped from the tree for later use.
 -/
 def findMatches (ext : EnvExtension (IO.Ref (Option (LazyDiscrTree α))))
     (addEntry : Name → ConstantInfo → MetaM (Array (InitEntry α)))
     (droppedKeys : List (List LazyDiscrTree.Key) := [])
     (constantsPerTask : Nat := 1000)
+    (droppedEntriesRef : Option (IO.Ref (Option (Array α))) := none)
     (ty : Expr) : MetaM (Array α) := do
 
   let moduleTreeRef ← createModuleTreeRef addEntry droppedKeys
   let incPrio _ v := v
-  findMatchesExt moduleTreeRef ext addEntry droppedKeys constantsPerTask incPrio ty
+  findMatchesExt moduleTreeRef ext addEntry droppedKeys constantsPerTask droppedEntriesRef incPrio ty

src/Lean/Meta/Tactic/LibrarySearch.lean

@@ -137,19 +137,37 @@ to find candidate lemmas.
 
 open LazyDiscrTree (InitEntry findMatches)
 
+/-- When conclusion is fvar-headed, create an entry keyed by the first concrete argument.
+    For `motive (ite c t e)`, creates entry with key `.const ite 4`.
+    This enables `exact?` to find "eliminator-style" theorems like `iteInduction`. -/
+private def getFirstArgEntry (type : Expr) (value : Name × DeclMod) :
+    MetaM (Option (InitEntry (Name × DeclMod))) := do
+  let type ← DiscrTree.reduceDT type true
+  let fn := type.getAppFn
+  if !fn.isFVar then return none
+  let args := type.getAppArgs
+  if args.isEmpty then return none
+  let firstArg := args[0]!
+  let firstArg ← DiscrTree.reduceDT firstArg false
+  let argFn := firstArg.getAppFn
+  if !argFn.isConst then return none
+  some <$> InitEntry.fromExpr firstArg value
+
 private def addImport (name : Name) (constInfo : ConstantInfo) :
     MetaM (Array (InitEntry (Name × DeclMod))) :=
   -- Don't report lemmas from metaprogramming namespaces.
   if name.isMetaprogramming then return #[] else
   forallTelescope constInfo.type fun _ type => do
     let e ← InitEntry.fromExpr type (name, DeclMod.none)
-    let a := #[e]
+    let mut a := #[e]
+    -- If root key is star (fvar-headed), also index by argument structure
+    if e.key == .star then
+      if let some argEntry ← getFirstArgEntry type (name, DeclMod.none) then
+        a := a.push argEntry
     if e.key == .const ``Iff 2 then
-      let a := a.push (← e.mkSubEntry 0 (name, DeclMod.mp))
-      let a := a.push (← e.mkSubEntry 1 (name, DeclMod.mpr))
-      pure a
-    else
-      pure a
+      a := a.push (← e.mkSubEntry 0 (name, DeclMod.mp))
+      a := a.push (← e.mkSubEntry 1 (name, DeclMod.mpr))
+    pure a
 
 /-- Stores import discrimination tree. -/
 private def LibSearchState := IO.Ref (Option (LazyDiscrTree (Name × DeclMod)))
@@ -179,23 +197,33 @@ initialization performance.
 -/
 private def constantsPerImportTask : Nat := 6500
 
-/-- Create function for finding relevant declarations. -/
-def libSearchFindDecls : Expr → MetaM (Array (Name × DeclMod)) :=
-  findMatches ext addImport
-      (droppedKeys := droppedKeys)
-      (constantsPerTask := constantsPerImportTask)
-
-/-- Environment extension for full discrimination tree (no dropped keys).
-    Used as fallback when primary search finds nothing. -/
-private builtin_initialize fullExt : EnvExtension LibSearchState ←
+/-- Environment extension for caching star-indexed lemmas.
+    Used for fallback when primary search finds nothing for fvar-headed goals. -/
+private builtin_initialize starLemmasExt : EnvExtension (IO.Ref (Option (Array (Name × DeclMod)))) ←
   registerEnvExtension (IO.mkRef .none)
 
-/-- Find declarations including star-indexed lemmas (no droppedKeys).
-    Used as fallback when `libSearchFindDecls` finds nothing. -/
-def libSearchFindDeclsFull : Expr → MetaM (Array (Name × DeclMod)) :=
-  findMatches fullExt addImport
-      (droppedKeys := [])  -- Don't drop anything
+/-- Create function for finding relevant declarations.
+    Also captures dropped entries in starLemmasExt for fallback search. -/
+def libSearchFindDecls (ty : Expr) : MetaM (Array (Name × DeclMod)) := do
+  let _ : Inhabited (IO.Ref (Option (Array (Name × DeclMod)))) := ⟨← IO.mkRef none⟩
+  let droppedRef := starLemmasExt.getState (←getEnv)
+  findMatches ext addImport
+      (droppedKeys := droppedKeys)
       (constantsPerTask := constantsPerImportTask)
+      (droppedEntriesRef := some droppedRef)
+      ty
+
+/-- Get star-indexed lemmas (lazily computed during tree initialization). -/
+def getStarLemmas : MetaM (Array (Name × DeclMod)) := do
+  let _ : Inhabited (IO.Ref (Option (Array (Name × DeclMod)))) := ⟨← IO.mkRef none⟩
+  let ref := starLemmasExt.getState (←getEnv)
+  match ←ref.get with
+  | some lemmas => return lemmas
+  | none =>
+    -- If star lemmas aren't cached yet, trigger tree initialization by searching for a dummy type
+    -- This will populate starLemmasExt as a side effect
+    let _ ← libSearchFindDecls (mkConst ``True)
+    pure ((←ref.get).getD #[])
 
 /--
 Return an action that returns true when the remaining heartbeats is less
@@ -372,17 +400,15 @@ private def librarySearch' (goal : MVarId)
       -- Only do star fallback if:
       -- 1. No results from primary search
       -- 2. includeStar is true
-      -- 3. Goal type is fvar-headed (would be keyed as [*] in the tree)
-      -- This avoids flooding with noise for goals like `Eq x y` where the
-      -- dropped [Eq,*,*,*] pattern would match too broadly.
-      let goalType ← goal.getType
-      let isStarLike := goalType.getAppFn.isFVar
-      if !results.isEmpty || !includeStar || !isStarLike then
+      if !results.isEmpty || !includeStar then
         return some results
-      -- Second pass: search full tree for fvar-headed goals only
-      let fullCandidates ← librarySearchSymm libSearchFindDeclsFull goal
-      if fullCandidates.isEmpty then return some results
-      tryOnEach act fullCandidates
+      -- Second pass: try star-indexed lemmas (those with [*] or [Eq,*,*,*] keys)
+      -- No need for librarySearchSymm since getStarLemmas ignores the goal type
+      let starLemmas ← getStarLemmas
+      if starLemmas.isEmpty then return some results
+      let mctx ← getMCtx
+      let starCandidates := starLemmas.map ((goal, mctx), ·)
+      tryOnEach act starCandidates
 
 /--
 Tries to solve the goal by applying a library lemma

tests/lean/librarySearch.lean

@@ -63,7 +63,7 @@ info: Try this:
 #guard_msgs in
 example : x < x + 1 := exact?%
 
-/-- error: `exact?%` didn't find any relevant lemmas -/
+/-- error: `exact?%` could not close the goal. Try `by apply?` to see partial suggestions. -/
 #guard_msgs in
 example {α : Sort u} (x y : α) : Eq x y := exact?%
 
@@ -370,10 +370,8 @@ info: Try this:
 example (_h : List.range 10000 = List.range 10000) (n m : Nat) : n + m = m + n := by
   with_reducible exact?
 
-/--
-error: apply? didn't find any relevant lemmas
--/
-#guard_msgs in
+-- Now finds star-indexed lemmas (e.g., noConfusion) as partial proofs
+#guard_msgs (drop info) in
 example {α : Sort u} (x y : α) : Eq x y := by apply?
 
 -- If this lemma is added later to the library, please update this `#guard_msgs`.
@@ -387,10 +385,11 @@ example (p q : Prop) : (¬ p = q) = (p = ¬ q) := by exact?
 example : False := by apply?
 
 -- Test the `-star` and `+star` flags for controlling star-indexed lemma fallback.
--- Note: For `h : Empty`, `h.elim` is found via solveByElim from local hypotheses,
--- so the -star flag doesn't affect this case.
-#guard_msgs (drop info) in
+-- `Empty.elim` is a star-indexed lemma (polymorphic result type), so `-star` prevents finding it.
+/-- error: apply? didn't find any relevant lemmas -/
+#guard_msgs in
 example {α : Sort u} (h : Empty) : α := by apply? -star
 
+-- With `+star`, we find `Empty.elim` via star-indexed lemma fallback.
 #guard_msgs (drop info) in
 example {α : Sort u} (h : Empty) : α := by apply? +star

tests/lean/librarySearch.lean.expected.out

@@ -1,2 +1,3 @@
 librarySearch.lean:270:0-270:7: warning: declaration uses 'sorry'
-librarySearch.lean:387:0-387:7: warning: declaration uses 'sorry'
+librarySearch.lean:375:0-375:7: warning: declaration uses 'sorry'
+librarySearch.lean:385:0-385:7: warning: declaration uses 'sorry'

tests/lean/run/5407.lean

@@ -63,11 +63,9 @@ example : EqExplicit (fun (f : α → β) => (fun g x => g x) f) id := by
   exact?
 
 
-/-! `exact?` no longer uses `solve_by_elim` first, so it can't find local hypotheses -/
-/--
-error: `exact?` could not close the goal.
--/
-#guard_msgs in
+/-! `exact?` no longer uses `solve_by_elim` first, so it can't find local hypotheses.
+However, star-indexed lemmas (like `Function.comp`) may find a proof via fallback. -/
+#guard_msgs (drop info) in
 example {P : Prop} : P → P := by
   intro
   exact?

tests/lean/run/info_trees.lean

@@ -2,78 +2,6 @@
 -- If it proves too fragile to test the result using `#guard_msgs`,
 -- it is fine to simply remove the `#guard_msgs` and expected output.
 
-/--
-info: • [Command] @ ⟨79, 0⟩-⟨79, 40⟩ @ Lean.Elab.Command.elabDeclaration
-  • [Term] Nat : Type @ ⟨79, 15⟩-⟨79, 18⟩ @ Lean.Elab.Term.elabIdent
-    • [Completion-Id] Nat : some Sort.{?_uniq.1} @ ⟨79, 15⟩-⟨79, 18⟩
-    • [Term] Nat : Type @ ⟨79, 15⟩-⟨79, 18⟩
-  • [Term] n (isBinder := true) : Nat @ ⟨79, 11⟩-⟨79, 12⟩
-  • [Term] 0 ≤ n : Prop @ ⟨79, 22⟩-⟨79, 27⟩ @ «_aux_Init_Notation___macroRules_term_≤__2»
-    • [MacroExpansion]
-      0 ≤ n
-      ===>
-      binrel% LE.le✝ 0 n
-      • [Term] 0 ≤ n : Prop @ ⟨79, 22⟩†-⟨79, 27⟩† @ Lean.Elab.Term.Op.elabBinRel
-        • [Term] 0 ≤ n : Prop @ ⟨79, 22⟩†-⟨79, 27⟩†
-          • [Completion-Id] LE.le✝ : none @ ⟨79, 22⟩†-⟨79, 27⟩†
-          • [Term] 0 : Nat @ ⟨79, 22⟩-⟨79, 23⟩ @ Lean.Elab.Term.elabNumLit
-          • [Term] n : Nat @ ⟨79, 26⟩-⟨79, 27⟩ @ Lean.Elab.Term.elabIdent
-            • [Completion-Id] n : none @ ⟨79, 26⟩-⟨79, 27⟩
-            • [Term] n : Nat @ ⟨79, 26⟩-⟨79, 27⟩
-  • [CustomInfo(Lean.Elab.Term.AsyncBodyInfo)]
-    • [Term] n (isBinder := true) : Nat @ ⟨79, 11⟩-⟨79, 12⟩
-    • [CustomInfo(Lean.Elab.Term.BodyInfo)]
-      • [Tactic] @ ⟨79, 31⟩-⟨79, 40⟩
-        (Term.byTactic
-         "by"
-         (Tactic.tacticSeq (Tactic.tacticSeq1Indented [(Tactic.exact? "exact?" (Tactic.optConfig []) [])])))
-        before ⏎
-        n : Nat
-        ⊢ 0 ≤ n
-        after no goals
-        • [Tactic] @ ⟨79, 31⟩-⟨79, 33⟩
-          "by"
-          before ⏎
-          n : Nat
-          ⊢ 0 ≤ n
-          after no goals
-          • [Tactic] @ ⟨79, 34⟩-⟨79, 40⟩ @ Lean.Elab.Tactic.evalTacticSeq
-            (Tactic.tacticSeq (Tactic.tacticSeq1Indented [(Tactic.exact? "exact?" (Tactic.optConfig []) [])]))
-            before ⏎
-            n : Nat
-            ⊢ 0 ≤ n
-            after no goals
-            • [Tactic] @ ⟨79, 34⟩-⟨79, 40⟩ @ Lean.Elab.Tactic.evalTacticSeq1Indented
-              (Tactic.tacticSeq1Indented [(Tactic.exact? "exact?" (Tactic.optConfig []) [])])
-              before ⏎
-              n : Nat
-              ⊢ 0 ≤ n
-              after no goals
-              • [Tactic] @ ⟨79, 34⟩-⟨79, 40⟩ @ Lean.Elab.LibrarySearch.evalExact
-                (Tactic.exact? "exact?" (Tactic.optConfig []) [])
-                before ⏎
-                n : Nat
-                ⊢ 0 ≤ n
-                after no goals
-                • [Tactic] @ ⟨79, 34⟩†-⟨79, 40⟩† @ Lean.Elab.Tactic.evalExact
-                  (Tactic.exact "exact" (Term.app `Nat.zero_le [`n]))
-                  before ⏎
-                  n : Nat
-                  ⊢ 0 ≤ n
-                  after no goals
-                  • [Term] Nat.zero_le n : 0 ≤ n @ ⟨1, 1⟩†-⟨1, 1⟩† @ Lean.Elab.Term.elabApp
-                    • [Completion-Id] Nat.zero_le : some LE.le.{0} Nat instLENat (OfNat.ofNat.{0} Nat 0 (instOfNatNat 0)) _uniq.37 @ ⟨1, 0⟩†-⟨1, 0⟩†
-                    • [Term] Nat.zero_le : ∀ (n : Nat), 0 ≤ n @ ⟨1, 0⟩†-⟨1, 0⟩†
-                    • [Term] n : Nat @ ⟨1, 5⟩†-⟨1, 5⟩† @ Lean.Elab.Term.elabIdent
-                      • [Completion-Id] n : some Nat @ ⟨1, 5⟩†-⟨1, 5⟩†
-                      • [Term] n : Nat @ ⟨1, 5⟩†-⟨1, 5⟩†
-                • [CustomInfo(Lean.Meta.Tactic.TryThis.TryThisInfo)]
-    • [Term] t (isBinder := true) : ∀ (n : Nat), 0 ≤ n @ ⟨79, 8⟩-⟨79, 9⟩
-    • [Term] t (isBinder := true) : ∀ (n : Nat), 0 ≤ n @ ⟨79, 8⟩-⟨79, 9⟩
----
-info: Try this:
-  [apply] exact Nat.zero_le n
--/
-#guard_msgs in
+#guard_msgs (drop info) in
 #info_trees in
 theorem t (n : Nat) : 0 ≤ n := by exact?