Repository: leanprover/lean4
Pull request #7842: feat: tuning and bugfix for MePo
Author: login: hanwenzhu
Labels: name: toolchain-available
description: A toolchain is available for this PR, at leanprover/lean4-pr-releases:pr-release-NNNN
name: changelog-language
description: Language features and metaprograms
Created: 2025-04-06T20:47:56Z | Merged: 2025-04-08T05:02:14Z

This PR to the `mepo` branch

* fixes the missing update to `relevant` in the function `Lean.PremiseSelection.MePo.mepo`
* adds a denylist for premise selection, which filters based on module name, premise name, and the resulting type, with filtered names stored in persistent environment extensions
* makes it possible for users to specify the `c` and `p` parameters to MePo at test time
* makes MePo conform to the `config.maxSuggestions` configuration.

(cc @kim-em)


Diff:
diff --git a/src/Lean/PremiseSelection/Basic.lean b/src/Lean/PremiseSelection/Basic.lean
index 272cc0f42236..55b3e1e789dc 100644
--- a/src/Lean/PremiseSelection/Basic.lean
+++ b/src/Lean/PremiseSelection/Basic.lean
@@ -69,6 +69,47 @@ structure Config where
 
 abbrev Selector : Type := MVarId → Config → MetaM (Array Suggestion)
 
+section DenyList
+
+/-- Premises from a module whose name has one of the following components are not retrieved. -/
+builtin_initialize moduleDenyListExt : SimplePersistentEnvExtension String (List String) ←
+  registerSimplePersistentEnvExtension {
+    addEntryFn := (·.cons)
+    addImportedFn := mkStateFromImportedEntries (·.cons) []
+  }
+
+/-- A premise whose name has one of the following components is not retrieved. -/
+builtin_initialize nameDenyListExt : SimplePersistentEnvExtension String (List String) ←
+  registerSimplePersistentEnvExtension {
+    addEntryFn := (·.cons)
+    addImportedFn := mkStateFromImportedEntries (·.cons) []
+  }
+
+/-- A premise whose `type.getForallBody.getAppFn` is a constant that has one of these prefixes is not retrieved. -/
+builtin_initialize typePrefixDenyListExt : SimplePersistentEnvExtension Name (List Name) ←
+  registerSimplePersistentEnvExtension {
+    addEntryFn := (·.cons)
+    addImportedFn := mkStateFromImportedEntries (·.cons) []
+  }
+
+def isDeniedModule (env : Environment) (moduleName : Name) : Bool :=
+  (moduleDenyListExt.getState env).any fun p => moduleName.anyS (· == p)
+
+def isDeniedPremise (env : Environment) (name : Name) : Bool := Id.run do
+  if name == ``sorryAx then return true
+  if name.isInternalDetail then return true
+  if (nameDenyListExt.getState env).any (fun p => name.anyS (· == p)) then return true
+  if let some moduleIdx := env.getModuleIdxFor? name then
+    let moduleName := env.header.moduleNames[moduleIdx.toNat]!
+    if isDeniedModule env moduleName then
+      return true
+  let some ci := env.find? name | return true
+  if let .const fnName _ := ci.type.getForallBody.getAppFn then
+    if (typePrefixDenyListExt.getState env).any (fun p => p.isPrefixOf fnName) then return true
+  return false
+
+end DenyList
+
 /--
 The trivial premise selector, which returns no suggestions.
 -/
@@ -84,7 +125,7 @@ def random (gen : StdGen := ⟨37, 59⟩) : Selector := fun _ cfg => do
   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
+    unless isDeniedPremise env name do
       suggestions := suggestions.push { name := name, score := 1.0 / consts.size.toFloat }
   return suggestions
 
diff --git a/src/Lean/PremiseSelection/MePo.lean b/src/Lean/PremiseSelection/MePo.lean
index 893fb395d05c..d00ed3435b82 100644
--- a/src/Lean/PremiseSelection/MePo.lean
+++ b/src/Lean/PremiseSelection/MePo.lean
@@ -43,8 +43,7 @@ def frequencyScore (frequency : Name → Nat) (relevant candidate : NameSet) : F
 
 def unweightedScore (relevant candidate : NameSet) : Float := weightedScore (fun _ => 1) relevant candidate
 
--- The values of p := 0.6 and c := 2.4 are taken from the MePo paper, and need to be tuned.
-def mepo (initialRelevant : NameSet) (score : NameSet → NameSet → Float) (accept : ConstantInfo → CoreM Bool) (p : Float := 0.6) (c : Float := 2.4) : CoreM (Array (Name × Float)) := do
+def mepo (initialRelevant : NameSet) (score : NameSet → NameSet → Float) (accept : ConstantInfo → CoreM Bool) (p : Float) (c : Float) : CoreM (Array (Name × Float)) := do
   let env ← getEnv
   let mut p := p
   let mut candidates := #[]
@@ -58,6 +57,7 @@ def mepo (initialRelevant : NameSet) (score : NameSet → NameSet → Float) (ac
     if newAccepted.isEmpty then return accepted
     accepted := newAccepted.foldl (fun acc (n, _, s) => acc.push (n, s)) accepted
     candidates := candidates'.map fun (n, c, _) => (n, c)
+    relevant := newAccepted.foldl (fun acc (_, ns, _) => acc ++ ns) relevant
     p := p + (1 - p) / c
   return accepted
 
@@ -72,7 +72,9 @@ end MePo
 
 open MePo
 
-def mepo (useRarity : Bool) (g : MVarId) : MetaM (Array Suggestion) := do
+-- The values of p := 0.6 and c := 2.4 are taken from the MePo paper, and need to be tuned.
+-- When retrieving ≤32 premises for use by downstream automation, Thomas Zhu suggests that c := 0.5 is optimal.
+def mepoSelector (useRarity : Bool) (p : Float := 0.6) (c : Float := 2.4) : Selector := fun g config => do
   let constants ← g.getConstants
   let env ← getEnv
   let score := if useRarity then
@@ -80,6 +82,12 @@ def mepo (useRarity : Bool) (g : MVarId) : MetaM (Array Suggestion) := do
     frequencyScore (frequency.findD · 0)
   else
     unweightedScore
-  let accept := fun ci => return ! (`Lean).isPrefixOf ci.name && Lean.Meta.allowCompletion env ci.name
-  let suggestions ← MePo.mepo constants score accept
-  return suggestions.map fun (n, s) => { name := n, score := s }
+  let accept := fun ci => do
+    return ! isDeniedPremise env ci.name
+  let suggestions ← mepo constants score accept p c
+  let suggestions := suggestions
+    |>.map (fun (n, s) => { name := n, score := s })
+    |>.reverse  -- we favor constants that appear at the end of `env.constants`
+  match config.maxSuggestions with
+  | none => return suggestions
+  | some k => return suggestions.take k
diff --git a/tests/lean/run/premise_selection_mepo.lean b/tests/lean/run/premise_selection_mepo.lean
index 24d2f8cf0a14..e70de1822586 100644
--- a/tests/lean/run/premise_selection_mepo.lean
+++ b/tests/lean/run/premise_selection_mepo.lean
@@ -1,6 +1,6 @@
 import Lean.PremiseSelection.MePo
 
-set_premise_selector fun g _ => Lean.PremiseSelection.mepo (useRarity := false) g
+set_premise_selector Lean.PremiseSelection.mepoSelector (useRarity := false)
 
 example (a b : Int) : a + b = b + a := by
   suggest_premises
@@ -11,7 +11,7 @@ example (x y z : List Int) : x ++ y ++ z = x ++ (y ++ z) := by
   suggest_premises
   sorry
 
-set_premise_selector fun g _ => Lean.PremiseSelection.mepo (useRarity := true) g
+set_premise_selector Lean.PremiseSelection.mepoSelector (useRarity := true)
 
 example (a b : Int) : a + b = b + a := by
   suggest_premises


Reviews:
user: login: kim-em
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-07T01:03:01Z
user: login: kim-em
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-07T01:03:16Z
user: login: kim-em
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-07T01:04:22Z
user: login: kim-em
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-07T01:04:48Z
user: login: hanwenzhu
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-07T01:35:27Z
user: login: hanwenzhu
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-07T01:38:48Z

Review comments:
diff_hunk: @@ -69,6 +69,54 @@ structure Config where
 
 abbrev Selector : Type := MVarId → Config → MetaM (Array Suggestion)
 
+section BlackList
+
+/-- Premises from a module whose name has one of the following components are not retrieved. -/
+builtin_initialize moduleBlackListExt : SimplePersistentEnvExtension String (List String) ←
+  let initialBlackList := [
+    "Aesop", "Auto", "Cli", "CodeAction", "DocGen4", "Duper", "ImportGraph", "Lake", "Lean",
+    "LeanSearchClient", "Linter", "Mathport", "MD4Lean", "Plausible", "ProofWidgets", "Qq",
+    "QuerySMT", "Tactic", "TacticExtra", "Test", "Testing", "UnicodeBasic", "Util"
+  ]
path: src/Lean/PremiseSelection/Basic.lean
user: login: kim-em
body: We can't have a hard coded list here that refers to downstream projects. This needs to be an extensible mechanism.
created_at: 2025-04-07T01:03:01Z
author_association: COLLABORATOR
diff_hunk: @@ -69,6 +69,54 @@ structure Config where
 
 abbrev Selector : Type := MVarId → Config → MetaM (Array Suggestion)
 
+section BlackList
+
+/-- Premises from a module whose name has one of the following components are not retrieved. -/
+builtin_initialize moduleBlackListExt : SimplePersistentEnvExtension String (List String) ←
+  let initialBlackList := [
path: src/Lean/PremiseSelection/Basic.lean
user: login: kim-em
body: Use `denyList` rather than `blackList`.
created_at: 2025-04-07T01:03:16Z
author_association: COLLABORATOR
diff_hunk: @@ -43,8 +43,9 @@ def frequencyScore (frequency : Name → Nat) (relevant candidate : NameSet) : F
 
 def unweightedScore (relevant candidate : NameSet) : Float := weightedScore (fun _ => 1) relevant candidate
 
--- The values of p := 0.6 and c := 2.4 are taken from the MePo paper, and need to be tuned.
-def mepo (initialRelevant : NameSet) (score : NameSet → NameSet → Float) (accept : ConstantInfo → CoreM Bool) (p : Float := 0.6) (c : Float := 2.4) : CoreM (Array (Name × Float)) := do
+-- The value c = 0.5 essentially forces the `while` loop to only iterate once, which preliminary
+-- experiments found to be optimal.
path: src/Lean/PremiseSelection/MePo.lean
user: login: kim-em
body: This seems really implausible to me, given the original MePo paper came up with 2.4. I'd want some pretty compelling evidence to back up this claim.
created_at: 2025-04-07T01:04:22Z
author_association: COLLABORATOR
diff_hunk: @@ -43,8 +43,9 @@ def frequencyScore (frequency : Name → Nat) (relevant candidate : NameSet) : F
 
 def unweightedScore (relevant candidate : NameSet) : Float := weightedScore (fun _ => 1) relevant candidate
 
--- The values of p := 0.6 and c := 2.4 are taken from the MePo paper, and need to be tuned.
-def mepo (initialRelevant : NameSet) (score : NameSet → NameSet → Float) (accept : ConstantInfo → CoreM Bool) (p : Float := 0.6) (c : Float := 2.4) : CoreM (Array (Name × Float)) := do
+-- The value c = 0.5 essentially forces the `while` loop to only iterate once, which preliminary
+-- experiments found to be optimal.
path: src/Lean/PremiseSelection/MePo.lean
user: login: kim-em
body: Let's back out this change for now until we can discuss benchmarking evidence?
created_at: 2025-04-07T01:04:48Z
author_association: COLLABORATOR
diff_hunk: @@ -69,6 +69,54 @@ structure Config where
 
 abbrev Selector : Type := MVarId → Config → MetaM (Array Suggestion)
 
+section BlackList
+
+/-- Premises from a module whose name has one of the following components are not retrieved. -/
+builtin_initialize moduleBlackListExt : SimplePersistentEnvExtension String (List String) ←
+  let initialBlackList := [
+    "Aesop", "Auto", "Cli", "CodeAction", "DocGen4", "Duper", "ImportGraph", "Lake", "Lean",
+    "LeanSearchClient", "Linter", "Mathport", "MD4Lean", "Plausible", "ProofWidgets", "Qq",
+    "QuerySMT", "Tactic", "TacticExtra", "Test", "Testing", "UnicodeBasic", "Util"
+  ]
path: src/Lean/PremiseSelection/Basic.lean
user: login: hanwenzhu
body: I am leaving the initial list empty, so for downstream use (e.g. Mathlib) at some place a deny list should be registered by ``modifyEnv fun env => moduleDenyListExt.addEntry env `Module``
created_at: 2025-04-07T01:35:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -69,6 +69,54 @@ structure Config where
 
 abbrev Selector : Type := MVarId → Config → MetaM (Array Suggestion)
 
+section BlackList
+
+/-- Premises from a module whose name has one of the following components are not retrieved. -/
+builtin_initialize moduleBlackListExt : SimplePersistentEnvExtension String (List String) ←
+  let initialBlackList := [
path: src/Lean/PremiseSelection/Basic.lean
user: login: hanwenzhu
body: Thanks! I changed this and also changed `isBlackListedPremise` to `isDeniedPremise`
created_at: 2025-04-07T01:38:48Z
author_association: CONTRIBUTOR

PR comments:
user: login: ghost
created_at: 2025-04-06T21:24:04Z
body: Mathlib CI status ([docs](https://leanprover-community.github.io/contribute/tags_and_branches.html)):
- ❗ Batteries/Mathlib CI will not be attempted unless your PR branches off the `nightly-with-mathlib` branch. Try `git rebase 8cefb2cf656f9e206df7dbc3910e73c597d6608b --onto da55b2e19b6af219f43f3599f1a72151799eb524`. You can force Mathlib CI using the `force-mathlib-ci` label. (2025-04-06 21:24:04)
author_association: NONE
user: login: hanwenzhu
created_at: 2025-04-07T19:46:39Z
body: changelog-language
author_association: CONTRIBUTOR