From 8d603d34dca8419a1f9ee3343d7edb520c5b8a40 Mon Sep 17 00:00:00 2001 From: Kim Morrison Date: Mon, 3 Nov 2025 03:55:34 +0100 Subject: [PATCH] feat: make `set_library_suggestions` persistent --- src/Lean/LibrarySuggestions.lean | 1 + src/Lean/LibrarySuggestions/Basic.lean | 94 ++++++++++++++++--- src/Lean/LibrarySuggestions/Default.lean | 25 +++++ tests/lean/run/library_suggestions.lean | 8 -- .../run/library_suggestions_override.lean | 38 ++++++++ .../run/library_suggestions_persistent.lean | 44 +++++++++ 6 files changed, 190 insertions(+), 20 deletions(-) create mode 100644 src/Lean/LibrarySuggestions/Default.lean create mode 100644 tests/lean/run/library_suggestions_override.lean create mode 100644 tests/lean/run/library_suggestions_persistent.lean diff --git a/src/Lean/LibrarySuggestions.lean b/src/Lean/LibrarySuggestions.lean index d187fb9d27..fcdeec0a39 100644 --- a/src/Lean/LibrarySuggestions.lean +++ b/src/Lean/LibrarySuggestions.lean @@ -10,3 +10,4 @@ import Lean.LibrarySuggestions.Basic import Lean.LibrarySuggestions.SymbolFrequency import Lean.LibrarySuggestions.MePo import Lean.LibrarySuggestions.SineQuaNon +import Lean.LibrarySuggestions.Default diff --git a/src/Lean/LibrarySuggestions/Basic.lean b/src/Lean/LibrarySuggestions/Basic.lean index f711c54775..f5eac90851 100644 --- a/src/Lean/LibrarySuggestions/Basic.lean +++ b/src/Lean/LibrarySuggestions/Basic.lean @@ -194,7 +194,7 @@ def maxSuggestions (selector : Selector) : Selector := fun g c => do return suggestions.take c.maxSuggestions /-- Combine two premise selectors, returning the best suggestions. -/ -def combine (selector₁ : Selector) (selector₂ : Selector) : Selector := fun g c => do +def combine (selector₁ selector₂ : Selector) : Selector := fun g c => do let suggestions₁ ← selector₁ g c let suggestions₂ ← selector₂ g c @@ -211,6 +211,50 @@ def combine (selector₁ : Selector) (selector₂ : Selector) : Selector := fun return sorted.take c.maxSuggestions +/-- +Combine two premise selectors by interspersing their results (ignoring scores). +The parameter `ratio` (defaulting to 0.5) controls the ratio of suggestions from each selector +while results are available from both. +-/ +def intersperse (selector₁ selector₂ : Selector) (ratio : Float := 0.5) : Selector := fun g c => do + -- Calculate how many suggestions to request from each selector based on the ratio + let max₁ := (c.maxSuggestions.toFloat * ratio).toUInt32.toNat + let max₂ := (c.maxSuggestions.toFloat * (1 - ratio)).toUInt32.toNat + + let suggestions₁ ← selector₁ g { c with maxSuggestions := max₁ } + let suggestions₂ ← selector₂ g { c with maxSuggestions := max₂ } + + let mut result := #[] + let mut i₁ := 0 + let mut i₂ := 0 + let mut count₁ := 0.0 + let mut count₂ := 0.0 + + -- Intersperse while both arrays have elements + while h : i₁ < suggestions₁.size ∧ i₂ < suggestions₂.size ∧ result.size < c.maxSuggestions do + -- Decide whether to take from selector₁ or selector₂ based on the ratio + let currentRatio := if count₁ + count₂ <= 0.0 then 0.0 else count₁ / (count₁ + count₂) + if currentRatio < ratio then + result := result.push suggestions₁[i₁] + i₁ := i₁ + 1 + count₁ := count₁ + 1 + else + result := result.push suggestions₂[i₂] + i₂ := i₂ + 1 + count₂ := count₂ + 1 + + -- Append remaining elements from selector₁ + while h : i₁ < suggestions₁.size ∧ result.size < c.maxSuggestions do + result := result.push suggestions₁[i₁] + i₁ := i₁ + 1 + + -- Append remaining elements from selector₂ + while h : i₂ < suggestions₂.size ∧ result.size < c.maxSuggestions do + result := result.push suggestions₂[i₂] + i₂ := i₂ + 1 + + return result + end Selector section DenyList @@ -302,12 +346,41 @@ def currentFile : Selector := fun _ cfg => do | _ => continue return suggestions -builtin_initialize librarySuggestionsExt : EnvExtension (Option Selector) ← - registerEnvExtension (pure none) +builtin_initialize librarySuggestionsExt : SimplePersistentEnvExtension Syntax (Option Syntax) ← + registerSimplePersistentEnvExtension { + addEntryFn := fun _ stx => some stx -- Last entry wins + addImportedFn := fun entries => + -- Take the last selector syntax from all imported modules + entries.foldl (init := none) fun acc moduleEntries => + moduleEntries.foldl (init := acc) fun _ stx => some stx + } + +/-- +Helper function to elaborate and evaluate selector syntax. +This is shared by both validation (`elabSetLibrarySuggestions`) and retrieval (`getSelector`). +-/ +def elabAndEvalSelector (stx : Syntax) : MetaM Selector := + Elab.Term.TermElabM.run' do + let selectorTerm ← Elab.Term.elabTermEnsuringType stx (some (Expr.const ``Selector [])) + unsafe Meta.evalExpr Selector (Expr.const ``Selector []) selectorTerm + +/-- +Get the currently registered library suggestions selector by evaluating the stored syntax. +Returns `none` if no selector is registered or if evaluation fails. + +Uses `Term.elabTermEnsuringType` to elaborate arbitrary syntax (not just identifiers). +-/ +def getSelector : MetaM (Option Selector) := do + let some stx := librarySuggestionsExt.getState (← getEnv) | return none + try + let selector ← elabAndEvalSelector stx + return some selector + catch _ => + return none /-- Generate library suggestions for the given metavariable, using the currently registered library suggestions engine. -/ def select (m : MVarId) (c : Config := {}) : MetaM (Array Suggestion) := do - let some selector := librarySuggestionsExt.getState (← getEnv) | + let some selector ← getSelector | throwError "No library suggestions engine registered. \ (Note that Lean does not provide a default library suggestions engine, \ these must be provided by a downstream library, \ @@ -324,23 +397,20 @@ library suggestions engines are configured via options in the `lakefile`, and commands are only used to override in a single declaration or file. -/ -/-- Set the current library suggestions engine.-/ -def registerLibrarySuggestions (selector : Selector) : CoreM Unit := do - modifyEnv fun env => librarySuggestionsExt.setState env (some selector) - open Lean Elab Command in @[builtin_command_elab setLibrarySuggestionsCmd, inherit_doc setLibrarySuggestionsCmd] def elabSetLibrarySuggestions : CommandElab | `(command| set_library_suggestions $selector) => do if `Lean.LibrarySuggestions.Basic ∉ (← getEnv).header.moduleNames then logWarning "Add `import Lean.LibrarySuggestions.Basic` before using the `set_library_suggestions` command." - let selector ← liftTermElabM do + -- Validate that the syntax can be elaborated (to catch errors early) + liftTermElabM do try - let selectorTerm ← Term.elabTermEnsuringType selector (some (Expr.const ``Selector [])) - unsafe Meta.evalExpr Selector (Expr.const ``Selector []) selectorTerm + discard <| elabAndEvalSelector selector catch _ => throwError "Failed to elaborate {selector} as a `MVarId → Config → MetaM (Array Suggestion)`." - liftCoreM (registerLibrarySuggestions selector) + -- Store the syntax (not the evaluated Selector) for persistence + modifyEnv fun env => librarySuggestionsExt.addEntry env selector | _ => throwUnsupportedSyntax open Lean.Elab.Tactic in diff --git a/src/Lean/LibrarySuggestions/Default.lean b/src/Lean/LibrarySuggestions/Default.lean new file mode 100644 index 0000000000..e9a192cfb1 --- /dev/null +++ b/src/Lean/LibrarySuggestions/Default.lean @@ -0,0 +1,25 @@ +/- +Copyright (c) 2025 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Kim Morrison +-/ +module + +prelude +import all Lean.LibrarySuggestions.SineQuaNon + +/-! +# Default library suggestions engine + +This module sets the default library suggestions engine to Sine Qua Non. +Any module that imports this (directly or transitively) will have library suggestions enabled. +-/ + +namespace Lean.LibrarySuggestions + +open Lean LibrarySuggestions SineQuaNon + +-- Set the default library suggestions engine to Sine Qua Non +set_library_suggestions Lean.LibrarySuggestions.sineQuaNonSelector + +end Lean.LibrarySuggestions diff --git a/tests/lean/run/library_suggestions.lean b/tests/lean/run/library_suggestions.lean index e3ab1b5cb6..71f64cab40 100644 --- a/tests/lean/run/library_suggestions.lean +++ b/tests/lean/run/library_suggestions.lean @@ -14,14 +14,6 @@ error: Failed to elaborate Nat as a `MVarId → Config → MetaM (Array Suggesti #guard_msgs in set_library_suggestions Nat -/-- -error: No library suggestions engine registered. (Note that Lean does not provide a default library suggestions engine, these must be provided by a downstream library, and configured using `set_library_suggestions`.) --/ -#guard_msgs in -example : True := by - suggestions - trivial - set_library_suggestions (fun _ _ => pure #[]) /-- info: Library suggestions: [] -/ diff --git a/tests/lean/run/library_suggestions_override.lean b/tests/lean/run/library_suggestions_override.lean new file mode 100644 index 0000000000..399df0f253 --- /dev/null +++ b/tests/lean/run/library_suggestions_override.lean @@ -0,0 +1,38 @@ +import Lean.LibrarySuggestions.Default + +/-! +# Test that set_library_suggestions can override the default + +This test verifies that even after importing Default (which sets Sine Qua Non as the default), +we can override it with a different selector. +-/ + +-- First, verify the default is set by using it +example {x : Dyadic} {prec : Int} : x.roundDown prec ≤ x := by + grind +suggestions + +-- Now override with an empty selector +set_library_suggestions (fun _ _ => pure #[]) + +-- Verify that grind +suggestions now fails because the empty selector provides no suggestions +/-- +error: `grind` failed +case grind +x : Dyadic +prec : Int +h : ¬x.roundDown prec ≤ x +⊢ False +[grind] Goal diagnostics + [facts] Asserted facts + [prop] ¬x.roundDown prec ≤ x + [eqc] False propositions + [prop] x.roundDown prec ≤ x + [cutsat] Assignment satisfying linear constraints + [assign] prec := 2 + [linarith] Linarith assignment for `Dyadic` + [assign] x := -1 + [assign] x.roundDown prec := 0 +-/ +#guard_msgs in +example {x : Dyadic} {prec : Int} : x.roundDown prec ≤ x := by + grind +suggestions diff --git a/tests/lean/run/library_suggestions_persistent.lean b/tests/lean/run/library_suggestions_persistent.lean new file mode 100644 index 0000000000..602a05ffb0 --- /dev/null +++ b/tests/lean/run/library_suggestions_persistent.lean @@ -0,0 +1,44 @@ +import Lean.LibrarySuggestions +import Lean.Meta.Basic +import Std.Data.ExtHashMap + +/-! +# Test that library suggestions persist across file boundaries + +This test verifies that the default library suggestion engine set in +`Lean.LibrarySuggestions.Default` is correctly persisted when imported via `Lean.LibrarySuggestions`. + +We do NOT call `set_library_suggestions` in this file - the selector should +already be set from importing Lean.LibrarySuggestions (which imports Default). +-/ + +/-- +info: ✓ Selector found in imported state: `Lean.LibrarySuggestions.sineQuaNonSelector +--- +info: ✓ Successfully retrieved selector using getSelector! +-/ +#guard_msgs in +open Lean Lean.LibrarySuggestions in +run_cmd do + let stx? := librarySuggestionsExt.getState (← getEnv) + match stx? with + | none => Lean.logInfo "❌ No selector found in imported state!" + | some stx => + Lean.logInfo s!"✓ Selector found in imported state: {stx}" + -- Try to retrieve the selector using getSelector + Elab.Command.liftTermElabM do + let selector? ← getSelector + match selector? with + | none => Lean.logInfo " ❌ getSelector returned none" + | some _ => Lean.logInfo " ✓ Successfully retrieved selector using getSelector!" + +-- These examples should work with grind +suggestions but not grind alone +-- (proving that the suggestions engine is active and helping) + +example {x : Dyadic} {prec : Int} : x.roundDown prec ≤ x := by + fail_if_success grind + grind +suggestions + +example {x : Dyadic} {prec : Int} : (x.roundUp prec).precision ≤ some prec := by + fail_if_success grind + grind +suggestions