From 8cefb2cf656f9e206df7dbc3910e73c597d6608b Mon Sep 17 00:00:00 2001 From: Kim Morrison Date: Fri, 14 Feb 2025 15:08:18 +1100 Subject: [PATCH] feat: premise selection API (#7061) This PR provides a basic API for a premise selection tool, which can be provided in downstream libraries. It does not implement premise selection itself! --- src/Init/Notation.lean | 7 ++ src/Init/Tactics.lean | 7 ++ src/Lean.lean | 1 + src/Lean/PremiseSelection.lean | 135 ++++++++++++++++++++++++++ tests/lean/run/premise_selection.lean | 38 ++++++++ 5 files changed, 188 insertions(+) create mode 100644 src/Lean/PremiseSelection.lean create mode 100644 tests/lean/run/premise_selection.lean diff --git a/src/Init/Notation.lean b/src/Init/Notation.lean index e4fd0a02a3..1d3d27702e 100644 --- a/src/Init/Notation.lean +++ b/src/Init/Notation.lean @@ -756,6 +756,13 @@ This is mostly useful for debugging info trees. syntax (name := infoTreesCmd) "#info_trees" " in" ppLine command : command +/-- +Specify a premise selection engine. +Note that Lean does not ship a default premise selection engine, +so this is only useful in conjunction with a downstream package which provides one. +-/ +syntax (name := setPremiseSelectorCmd) + "set_premise_selector" term : command namespace Parser diff --git a/src/Init/Tactics.lean b/src/Init/Tactics.lean index db1bdb9009..45ef2c8e93 100644 --- a/src/Init/Tactics.lean +++ b/src/Init/Tactics.lean @@ -1589,6 +1589,13 @@ as well as tactics such as `next`, `case`, and `rename_i`. -/ syntax (name := exposeNames) "expose_names" : tactic +/-- +`#suggest_premises` will suggest premises for the current goal, using the currently registered premise selector. + +The suggestions are printed in the order of their confidence, from highest to lowest. +-/ +syntax (name := suggestPremises) "suggest_premises" : tactic + /-- Close fixed-width `BitVec` and `Bool` goals by obtaining a proof from an external SAT solver and verifying it inside Lean. The solvable goals are currently limited to diff --git a/src/Lean.lean b/src/Lean.lean index ee13853210..ecdc9393f7 100644 --- a/src/Lean.lean +++ b/src/Lean.lean @@ -38,3 +38,4 @@ import Lean.LabelAttribute import Lean.AddDecl import Lean.Replay import Lean.PrivateName +import Lean.PremiseSelection diff --git a/src/Lean/PremiseSelection.lean b/src/Lean/PremiseSelection.lean new file mode 100644 index 0000000000..272cc0f422 --- /dev/null +++ b/src/Lean/PremiseSelection.lean @@ -0,0 +1,135 @@ +/- +Copyright (c) 2025 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Kim Morrison +-/ +prelude +import Lean.Elab.Command +import Lean.Meta.Eval +import Lean.Meta.CompletionName +import Init.Data.Random + +/-! +# An API for premise selection algorithms. + +This module provides a basic API for premise selection algorithms, +which are used to suggest identifiers that should be introduced in a proof. + +The core interface is the `Selector` type, which is a function from a metavariable +and a configuration to a list of suggestions. + +The `Selector` is registered as an environment extension, and the trivial (no suggestions) implementation +is `Lean.PremiseSelection.empty`. + +Lean does not provide a default premise selector, so this module is intended to be used in conjunction +with a downstream package which registers a premise selector. +-/ + +namespace Lean.PremiseSelection + +/-- +A `Suggestion` is essentially just an identifier and a confidence score that the identifier is relevant. +If the premise selection request included information about the intended use (e.g. in the simplifier, in `grind`, etc.) +the score may be adjusted for that application. + +-/ +structure Suggestion where + name : Name + /-- + The score of the suggestion, as a probability that this suggestion should be used. + -/ + score : Float + +structure Config where + /-- + The maximum number of suggestions to return. + -/ + maxSuggestions : Option Nat := none + /-- + The tactic that is calling the premise selection, e.g. `simp`, `grind`, or `aesop`. + This may be used to adjust the score of the suggestions + -/ + caller : Option Name := none + /-- + A filter on suggestions; only suggestions returning `true` should be returned. + (It can be better to filter on the premise selection side, to ensure that enough suggestions are returned.) + -/ + filter : Name → MetaM Bool := fun _ => pure true + /-- + An optional arbitrary "hint" to the premise selection algorithm. + There is no guarantee that the algorithm will make any use of the hint. + + Potential use cases include a natural language comment provided by the user + (e.g. allowing use of the premise selector as a search engine) + or including context from the current proof and/or file. + + We may later split these use cases into separate fields if necessary. + -/ + hint : Option String := none + +abbrev Selector : Type := MVarId → Config → MetaM (Array Suggestion) + +/-- +The trivial premise selector, which returns no suggestions. +-/ +def empty : Selector := fun _ _ => pure #[] + +/-- A random premise selection algorithm, provided solely for testing purposes. -/ +def random (gen : StdGen := ⟨37, 59⟩) : Selector := fun _ cfg => do + IO.stdGenRef.set gen + let env ← getEnv + let max := cfg.maxSuggestions.getD 10 + let consts := env.const2ModIdx.keysArray + let mut suggestions := #[] + while suggestions.size < max do + let i ← IO.rand 0 consts.size + let name := consts[i]! + if ! (`Lean).isPrefixOf name && Lean.Meta.allowCompletion env name then + suggestions := suggestions.push { name := name, score := 1.0 / consts.size.toFloat } + return suggestions + +initialize premiseSelectorExt : EnvExtension (Option Selector) ← + registerEnvExtension (pure none) + +/-- Generate premise suggestions for the given metavariable, using the currently registered premise selector. -/ +def select (m : MVarId) (c : Config := {}) : MetaM (Array Suggestion) := do + let some selector := premiseSelectorExt.getState (← getEnv) | + throwError "No premise selector registered. \ + (Note the Lean does not provide a default premise selector, these must be installed by a downstream library.)" + selector m c + +/-! +Currently the registration mechanism is just global state. +This means that if multiple modules register premise selectors, +the behaviour will be dependent on the order of loading modules. + +We should replace this with a mechanism so that +premise selectors are configured via options in the `lakefile`, and +commands are only used to override in a single declaration or file. +-/ + +/-- Set the current premise selector.-/ +def registerPremiseSelector (selector : Selector) : CoreM Unit := do + modifyEnv fun env => premiseSelectorExt.setState env (some selector) + +open Lean Elab Command in +@[builtin_command_elab setPremiseSelectorCmd, inherit_doc setPremiseSelectorCmd] +def elabSetPremiseSelector : CommandElab + | `(command| set_premise_selector $selector) => do + let selector ← liftTermElabM do + try + let selectorTerm ← Term.elabTermEnsuringType selector (some (Expr.const ``Selector [])) + unsafe Meta.evalExpr Selector (Expr.const ``Selector []) selectorTerm + catch _ => + throwError "Failed to elaborate {selector} as a `MVarId → Config → MetaM (Array Suggestion)`." + liftCoreM (registerPremiseSelector selector) + | _ => throwUnsupportedSyntax + +open Lean.Elab.Tactic in +@[builtin_tactic Lean.Parser.Tactic.suggestPremises] def evalSuggestPremises : Tactic := fun _ => + liftMetaTactic1 fun mvarId => do + let suggestions ← select mvarId + logInfo m!"Premise suggestions: {suggestions.map (·.name)}" + return mvarId + +end Lean.PremiseSelection diff --git a/tests/lean/run/premise_selection.lean b/tests/lean/run/premise_selection.lean new file mode 100644 index 0000000000..da0f006e2a --- /dev/null +++ b/tests/lean/run/premise_selection.lean @@ -0,0 +1,38 @@ +import Lean.PremiseSelection + +/-- +error: type mismatch + Nat +has type + Type : Type 1 +but is expected to have type + Lean.PremiseSelection.Selector : Type +--- +error: Failed to elaborate Nat as a `MVarId → Config → MetaM (Array Suggestion)`. +-/ +#guard_msgs in +set_premise_selector Nat + +/-- +error: No premise selector registered. (Note the Lean does not provide a default premise selector, these must be installed by a downstream library.) +-/ +#guard_msgs in +example : True := by + suggest_premises + trivial + +set_premise_selector (fun _ _ => pure #[]) + +/-- info: Premise suggestions: [] -/ +#guard_msgs in +example : True := by + suggest_premises + trivial + +set_premise_selector Lean.PremiseSelection.random ⟨1,1⟩ + +-- This would be an extremely fragile test (select 10 random constants!) +-- so we do not use #guard_msgs. +example : True := by + suggest_premises + trivial