feat: ematch theorem activation

src/Lean/Meta/Tactic/Grind/Internalize.lean

@@ -37,6 +37,21 @@ private def updateAppMap (e : Expr) : GoalM Unit := do
       s.appMap.insert key [e]
   }
 
+private def activateTheoremPatterns (fName : Name) : GoalM Unit := do
+  if let some thms := (← get).thmMap.find? fName then
+    modify fun s => { s with thmMap := s.thmMap.erase fName }
+    let appMap := (← get).appMap
+    for thm in thms do
+      let symbols := thm.symbols.filter fun sym => !appMap.contains sym
+      let thm := { thm with symbols }
+      match symbols with
+      | [] =>
+        trace[grind.pattern] "activated `{thm.origin.key}`"
+        modify fun s => { s with newThms := s.newThms.push thm }
+      | _ =>
+        trace[grind.pattern] "reinsert `{thm.origin.key}`"
+        modify fun s => { s with thmMap := s.thmMap.insertTheorem thm }
+
 partial def internalize (e : Expr) (generation : Nat) : GoalM Unit := do
   if (← alreadyInternalized e) then return ()
   match e with
@@ -63,7 +78,9 @@ partial def internalize (e : Expr) (generation : Nat) : GoalM Unit := do
         internalize c generation
         registerParent e c
       else
-        unless f.isConst do
+        if let .const fName _ := f then
+          activateTheoremPatterns fName
+        else
           internalize f generation
           registerParent e f
         for h : i in [: args.size] do

src/Lean/Meta/Tactic/Grind/Run.lean

@@ -46,7 +46,8 @@ def GrindM.run (x : GrindM α) (mainDeclName : Name) : MetaM α := do
 def mkGoal (mvarId : MVarId) : GrindM Goal := do
   let trueExpr ← getTrueExpr
   let falseExpr ← getFalseExpr
-  GoalM.run' { mvarId } do
+  let thmMap ← getTheoremPatterns
+  GoalM.run' { mvarId, thmMap } do
     mkENodeCore falseExpr (interpreted := true) (ctor := false) (generation := 0)
     mkENodeCore trueExpr (interpreted := true) (ctor := false) (generation := 0)
 

src/Lean/Meta/Tactic/Grind/TheoremPatterns.lean

@@ -25,6 +25,14 @@ inductive Origin where
   | other
   deriving Inhabited, Repr
 
+/-- A unique identifier corresponding to the origin. -/
+def Origin.key : Origin → Name
+  | .decl declName => declName
+  | .fvar fvarId   => fvarId.name
+  | .stx id _      => id
+  | .other         => `other
+
+-- TODO: find better name
 structure TheoremPattern where
   proof       : Expr
   numParams   : Nat
@@ -34,16 +42,21 @@ structure TheoremPattern where
   origin      : Origin
   deriving Inhabited
 
-abbrev TheoremPatterns := SMap Name (List TheoremPattern)
+-- TODO: find better name
+abbrev TheoremPatterns := PHashMap Name (List TheoremPattern)
+
+def TheoremPatterns.insertTheorem (s : TheoremPatterns) (thm : TheoremPattern) : TheoremPatterns := Id.run do
+  let .const declName :: syms := thm.symbols
+    | unreachable!
+  let thm := { thm with symbols := syms }
+  if let some thms := s.find? declName then
+    return s.insert declName (thm::thms)
+  else
+    return s.insert declName [thm]
 
 builtin_initialize theoremPatternsExt : SimpleScopedEnvExtension TheoremPattern TheoremPatterns ←
   registerSimpleScopedEnvExtension {
-    addEntry := fun s t => Id.run do
-      let .const declName :: _ := t.symbols | unreachable!
-      if let some ts := s.find? declName then
-        s.insert declName (t::ts)
-      else
-        s.insert declName [t]
+    addEntry := TheoremPatterns.insertTheorem
     initial  := .empty
   }
 
@@ -298,4 +311,7 @@ def addTheoremPattern (declName : Name) (numParams : Nat) (patterns : List Expr)
      origin := .decl declName
   }
 
+def getTheoremPatterns : CoreM TheoremPatterns :=
+  return theoremPatternsExt.getState (← getEnv)
+
 end Lean.Meta.Grind

src/Lean/Meta/Tactic/Grind/Types.lean

@@ -13,6 +13,7 @@ import Lean.Meta.Tactic.Simp.Types
 import Lean.Meta.Tactic.Util
 import Lean.Meta.Tactic.Grind.Canon
 import Lean.Meta.Tactic.Grind.Attr
+import Lean.Meta.Tactic.Grind.TheoremPatterns
 
 namespace Lean.Meta.Grind
 
@@ -273,6 +274,15 @@ structure Goal where
   gmt          : Nat := 0
   /-- Next unique index for creating ENodes -/
   nextIdx      : Nat := 0
+  /-- Active theorems that we have performed ematching at least once. -/
+  thms         : PArray TheoremPattern := {}
+  /-- Active theorems that we have not performed any round of ematching yet. -/
+  newThms      : PArray TheoremPattern := {}
+  /--
+  Inactive global theorems. As we internalize terms, we activate theorems as we find their symbols.
+  Local theorem provided by users are added directly into `newThms`.
+  -/
+  thmMap       : TheoremPatterns
   deriving Inhabited
 
 def Goal.admit (goal : Goal) : MetaM Unit :=

tests/lean/run/grind_pattern2.lean

@@ -0,0 +1,39 @@
+def Set (α : Type) := α → Bool
+
+def insertElem [DecidableEq α] (s : Set α) (a : α) : Set α :=
+  fun x => a = x || s x
+
+def contains (s : Set α) (a : α) : Bool :=
+  s a
+
+theorem contains_insert [DecidableEq α] (s : Set α) (a : α) : contains (insertElem s a) a := by
+  simp [contains, insertElem]
+
+grind_pattern contains_insert => contains (insertElem s a) a
+
+-- TheoremPattern activation test
+
+set_option trace.grind.pattern true
+
+/--
+warning: declaration uses 'sorry'
+---
+info: [grind.pattern] activated `contains_insert`
+-/
+#guard_msgs in
+example [DecidableEq α] (s₁ s₂ : Set α) (a₁ a₂ : α) :
+        s₂ = insertElem s₁ a₁ → a₁ = a₂ → contains s₂ a₂ := by
+  fail_if_success grind
+  sorry
+
+/--
+warning: declaration uses 'sorry'
+---
+info: [grind.pattern] reinsert `contains_insert`
+[grind.pattern] activated `contains_insert`
+-/
+#guard_msgs in
+example [DecidableEq α] (s₁ s₂ : Set α) (a₁ a₂ : α) :
+        ¬ contains s₂ a₂ → s₂ = insertElem s₁ a₁ → a₁ = a₂ → False := by
+  fail_if_success grind
+  sorry