feat: make set_library_suggestions persistent

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

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

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

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: [] -/

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

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