feat: include symbols in ground grind patterns (#11589)

This PR improves indexing for `grind` patterns. We now include symbols
occurring in nested ground patterns. This important to minimize the
number of activated E-match theorems.

Author: leodemoura

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

@@ -7,6 +7,7 @@ module
 prelude
 public import Lean.Meta.Tactic.Grind.Theorems
 import Init.Grind.Util
+import Lean.Util.ForEachExpr
 import Lean.Meta.Tactic.Grind.Util
 import Lean.Meta.Match.Basic
 import Lean.Meta.Tactic.TryThis
@@ -580,6 +581,22 @@ private def saveSymbol (h : HeadIndex) : M Unit := do
   unless (← get).symbolSet.contains h do
     modify fun s => { s with symbols := s.symbols.push h, symbolSet := s.symbolSet.insert h }
 
+private def saveSymbolsAt (e : Expr) : M Unit := do
+  e.forEach' fun e => do
+    if e.isApp || e.isConst then
+      /- **Note**: We ignore function symbols that have special handling in the internalizer. -/
+      if let .const declName _ := e.getAppFn then
+        if declName == ``OfNat.ofNat || declName == ``Grind.nestedProof
+           || declName == ``Grind.eqBwdPattern
+           || declName == ``Grind.nestedDecidable || declName == ``ite then
+          return false
+    match e with
+    | .const .. =>
+      saveSymbol e.toHeadIndex
+      return false
+    | _ =>
+      return true
+
 private def foundBVar (idx : Nat) : M Bool :=
   return (← get).bvarsFound.contains idx
 
@@ -672,30 +689,48 @@ private def getPatternFn? (pattern : Expr) (inSupport : Bool) (root : Bool) (arg
 
 private partial def go (pattern : Expr) (inSupport : Bool) (root : Bool) : M Expr := do
   if let some (e, k) := isOffsetPattern? pattern then
-    let e ← goArg e inSupport .relevant
+    let e ← goArg e inSupport .relevant (isEqBwdParent := false)
     if e == dontCare then
       return dontCare
     else
       return mkOffsetPattern e k
   let some f ← getPatternFn? pattern inSupport root .relevant
     | throwError "invalid pattern, (non-forbidden) application expected{indentD (ppPattern pattern)}"
   assert! f.isConst || f.isFVar
-  unless f.isConstOf ``Grind.eqBwdPattern do
+  let isEqBwd := f.isConstOf ``Grind.eqBwdPattern
+  unless isEqBwd do
     saveSymbol f.toHeadIndex
   let mut args := pattern.getAppArgs.toVector
   let patternArgKinds ← getPatternArgKinds f args.size
   for h : i in *...args.size do
     let arg := args[i]
     let argKind := patternArgKinds[i]?.getD .relevant
-    args := args.set i (← goArg arg (inSupport || argKind.isSupport) argKind)
+    args := args.set i (← goArg arg (inSupport || argKind.isSupport) argKind isEqBwd)
   return mkAppN f args.toArray
 where
-  goArg (arg : Expr) (inSupport : Bool) (argKind : PatternArgKind) : M Expr := do
+  goArg (arg : Expr) (inSupport : Bool) (argKind : PatternArgKind) (isEqBwdParent : Bool) : M Expr := do
     if !arg.hasLooseBVars then
       if arg.hasMVar then
         pure dontCare
+      else if (← isProof arg) then
+        pure dontCare
       else
-        return mkGroundPattern (← expandOffsetPatterns arg)
+        let arg ← expandOffsetPatterns arg
+        unless isEqBwdParent do
+          /-
+          **Note**: We ignore symbols in ground patterns if the parent is the auxiliary ``Grind.eqBwdPattern
+          We do that because we want to sign an error in examples such as:
+          ```
+          theorem dummy (x : Nat) : x = x :=
+            rfl
+          -- error: invalid pattern for `dummy`
+          --  [@Lean.Grind.eqBwdPattern `[Nat] #0 #0]
+          -- the pattern does not contain constant symbols for indexing
+          attribute [grind ←=] dummy
+          ```
+          -/
+          saveSymbolsAt arg
+        return mkGroundPattern arg
     else match arg with
       | .bvar idx =>
         if inSupport && (← foundBVar idx) then

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

@@ -89,12 +89,17 @@ where
     unless (← isEqFalse e) do return false
     return !(← isEqFalse e')
 
+def updateIndicesFound (k : HeadIndex) : GoalM Unit := do
+  if (← get).indicesFound.contains k then return ()
+  modify fun s => { s with indicesFound := s.indicesFound.insert k }
+
 /--
 Given an application `e` of the form `f a_1 ... a_n`,
 adds entry `f ↦ e` to `appMap`. Recall that `appMap` is a multi-map.
 -/
 private def updateAppMap (e : Expr) : GoalM Unit := do
   let key := e.toHeadIndex
+  updateIndicesFound key
   trace_goal[grind.debug.appMap] "{e} => {repr key}"
   modify fun s => { s with
     appMap := if let some es := s.appMap.find? key then
@@ -280,10 +285,10 @@ private def activateTheoremsCore [TheoremLike α] (declName : Name)
       let origin := TheoremLike.getOrigin thm
       trace_goal[grind.debug.theorem.activate] "`{declName}` => `{origin.key}`"
       unless s.isErased origin do
-        let appMap  := (← get).appMap
-        let symbols := TheoremLike.getSymbols thm
-        let symbols := symbols.filter fun sym => !appMap.contains sym
-        let thm     := TheoremLike.setSymbols thm symbols
+        let indicesFound := (← get).indicesFound
+        let symbols      := TheoremLike.getSymbols thm
+        let symbols      := symbols.filter fun sym => !indicesFound.contains sym
+        let thm          := TheoremLike.setSymbols thm symbols
         match symbols with
         | [] =>
           trace_goal[grind.debug.theorem.activate] "`{origin.key}`"
@@ -515,6 +520,7 @@ where
     | .lit .. =>
       mkENode e generation
     | .const declName _ =>
+      updateIndicesFound (.const declName)
       mkENode e generation
       activateTheorems declName generation
     | .mvar .. =>

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

@@ -121,6 +121,11 @@ def assertExtra (params : Params) : GoalM Unit := do
   for thm in params.extraInj do
     activateInjectiveTheorem thm 0
 
+private def initENodeCore (e : Expr) (interpreted ctor : Bool) : GoalM Unit := do
+  if let .const declName _ := e then
+    updateIndicesFound (.const declName)
+  mkENodeCore e interpreted ctor (generation := 0) (funCC := false)
+
 private def mkGoal (mvarId : MVarId) (params : Params) : GrindM Goal := do
   let mvarId ← if params.config.clean then mvarId.exposeNames else pure mvarId
   let trueExpr ← getTrueExpr
@@ -134,12 +139,12 @@ private def mkGoal (mvarId : MVarId) (params : Params) : GrindM Goal := do
   let clean ← mkCleanState mvarId params
   let sstates ← Solvers.mkInitialStates
   GoalM.run' { mvarId, ematch.thmMap := thmMap, inj.thms := params.inj, split.casesTypes := casesTypes, clean, sstates } do
-    mkENodeCore falseExpr (interpreted := true) (ctor := false) (generation := 0) (funCC := false)
-    mkENodeCore trueExpr (interpreted := true) (ctor := false) (generation := 0) (funCC := false)
-    mkENodeCore btrueExpr (interpreted := false) (ctor := true) (generation := 0) (funCC := false)
-    mkENodeCore bfalseExpr (interpreted := false) (ctor := true) (generation := 0) (funCC := false)
-    mkENodeCore natZeroExpr (interpreted := true) (ctor := false) (generation := 0) (funCC := false)
-    mkENodeCore ordEqExpr (interpreted := false) (ctor := true) (generation := 0) (funCC := false)
+    initENodeCore falseExpr (interpreted := true) (ctor := false)
+    initENodeCore trueExpr (interpreted := true) (ctor := false)
+    initENodeCore btrueExpr (interpreted := false) (ctor := true)
+    initENodeCore bfalseExpr (interpreted := false) (ctor := true)
+    initENodeCore natZeroExpr (interpreted := true) (ctor := false)
+    initENodeCore ordEqExpr (interpreted := false) (ctor := true)
     assertExtra params
 
 structure Result where

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

@@ -951,6 +951,11 @@ structure Goal where
   it is its unique id.
   -/
   appMap       : PHashMap HeadIndex (List Expr) := {}
+  /--
+  All constants (*not* in `appMap`) that have been internalized, *and*
+  `appMap`'s domain. We use this collection during theorem activation.
+  -/
+  indicesFound : PHashSet HeadIndex := {}
   /-- Equations and propositions to be processed. -/
   newFacts     : Array NewFact := #[]
   /-- `inconsistent := true` if `ENode`s for `True` and `False` are in the same equivalence class. -/

stage0/src/stdlib_flags.h

@@ -1,3 +1,4 @@
+// update me!
 #include "util/options.h"
 
 namespace lean {

tests/lean/run/grind_const_pattern.lean

@@ -73,11 +73,8 @@ h : ¬f x = 11
     [prop] ¬f x = 11
   [eqc] False propositions
     [prop] f x = 11
-  [ematch] E-matching patterns
-    [thm] fa: [f `[a]]
   [cutsat] Assignment satisfying linear constraints
     [assign] x := 1
-    [assign] a := 3
     [assign] f x := 2
 -/
 #guard_msgs (error) in

tests/lean/run/grind_pattern2.lean

@@ -45,13 +45,6 @@ trace: [grind.internalize] [0] x
 [grind.internalize] [0] y
 [grind.internalize] [0] z
 [grind.internalize] [0] foo x y
-[grind.internalize] [0] [a, b]
-[grind.internalize] [0] Nat
-[grind.internalize] [0] a
-[grind.internalize] [0] [b]
-[grind.internalize] [0] b
-[grind.internalize] [0] []
-[grind.ematch] activated `fooThm`, [foo #0 `[[a, b]]]
 -/
 #guard_msgs (trace) in
 set_option trace.grind.internalize true in

tests/lean/run/try_induction.lean

@@ -130,14 +130,14 @@ theorem hyperoperation_recursion (n m k : ℕ) :
 /--
 info: Try these:
   [apply] (induction k) <;> grind
-  [apply] (induction k) <;> grind only [hyperoperation, = add_zero, = add_succ, = hyperoperation_zero]
+  [apply] (induction k) <;> grind only [hyperoperation, = add_zero, = add_succ]
   [apply] ·
     induction k
     · grind => instantiate only [hyperoperation, = add_zero]
     ·
       grind =>
         instantiate only [hyperoperation, = add_succ]
-        instantiate only [= hyperoperation_zero]
+        instantiate only [hyperoperation]
 -/
 #guard_msgs in
 @[grind =]