fix: E-matching thresholds in the grind tactic (#6536)

This PR fixes different thresholds for controlling E-matching in the
`grind` tactic.

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

@@ -137,7 +137,8 @@ private partial def processMatch (c : Choice) (p : Expr) (e : Expr) : M Unit :=
   let mut curr := e
   repeat
     let n ← getENode curr
-    if n.generation <= maxGeneration
+    -- Remark: we use `<` because the instance generation is the maximum term generation + 1
+    if n.generation < maxGeneration
        -- uses heterogeneous equality or is the root of its congruence class
        && (n.heqProofs || n.isCongrRoot)
        && eqvFunctions pFn curr.getAppFn
@@ -155,7 +156,7 @@ private partial def processOffset (c : Choice) (pArg : Expr) (k : Nat) (e : Expr
   let mut curr := e
   repeat
     let n ← getENode curr
-    if n.generation <= maxGeneration then
+    if n.generation < maxGeneration then
       if let some (eArg, k') ← isOffset? curr |>.run then
         if k' < k then
           let c := c.updateGen n.generation
@@ -166,14 +167,14 @@ private partial def processOffset (c : Choice) (pArg : Expr) (k : Nat) (e : Expr
         else if k' > k then
           let eArg' := mkNatAdd eArg (mkNatLit (k' - k))
           let eArg' ← shareCommon (← canon eArg')
-          internalize eArg' n.generation
+          internalize eArg' (n.generation+1)
           if let some c ← matchArg? c pArg eArg' |>.run then
             pushChoice (c.updateGen n.generation)
       else if let some k' ← evalNat curr |>.run then
         if k' >= k then
           let eArg' := mkNatLit (k' - k)
           let eArg' ← shareCommon (← canon eArg')
-          internalize eArg' n.generation
+          internalize eArg' (n.generation+1)
           if let some c ← matchArg? c pArg eArg' |>.run then
             pushChoice (c.updateGen n.generation)
     curr ← getNext curr
@@ -186,7 +187,7 @@ private def processContinue (c : Choice) (p : Expr) : M Unit := do
   let maxGeneration ← getMaxGeneration
   for app in apps do
     let n ← getENode app
-    if n.generation <= maxGeneration
+    if n.generation < maxGeneration
        && (n.heqProofs || n.isCongrRoot) then
       if let some c ← matchArgs? c p app |>.run then
         let gen := n.generation
@@ -214,7 +215,7 @@ private def addNewInstance (origin : Origin) (proof : Expr) (generation : Nat) :
   if Match.isMatchEqnTheorem (← getEnv) origin.key then
     prop ← annotateMatchEqnType prop
   trace[grind.ematch.instance] "{← origin.pp}: {prop}"
-  addTheoremInstance proof prop generation
+  addTheoremInstance proof prop (generation+1)
 
 /--
 After processing a (multi-)pattern, use the choice assignment to instantiate the proof.
@@ -262,17 +263,18 @@ where
     isDefEq x val
 
 /-- Process choice stack until we don't have more choices to be processed. -/
-private partial def processChoices : M Unit := do
-  unless (← get).choiceStack.isEmpty do
+private def processChoices : M Unit := do
+  let maxGeneration ← getMaxGeneration
+  while !(← get).choiceStack.isEmpty do
     checkSystem "ematch"
     if (← checkMaxInstancesExceeded) then return ()
     let c ← modifyGet fun s : State => (s.choiceStack.head!, { s with choiceStack := s.choiceStack.tail! })
-    match c.cnstrs with
-    | [] => instantiateTheorem c
-    | .match p e :: cnstrs => processMatch { c with cnstrs } p e
-    | .offset p k e :: cnstrs => processOffset { c with cnstrs } p k e
-    | .continue p :: cnstrs => processContinue { c with cnstrs } p
-    processChoices
+    if c.gen < maxGeneration then
+      match c.cnstrs with
+      | [] => instantiateTheorem c
+      | .match p e :: cnstrs => processMatch { c with cnstrs } p e
+      | .offset p k e :: cnstrs => processOffset { c with cnstrs } p k e
+      | .continue p :: cnstrs => processContinue { c with cnstrs } p
 
 private def main (p : Expr) (cnstrs : List Cnstr) : M Unit := do
   let some apps := (← getThe Goal).appMap.find? p.toHeadIndex
@@ -327,7 +329,7 @@ def ematch : GoalM Unit := do
   let go (thms newThms : PArray EMatchTheorem) : EMatch.M Unit := do
     withReader (fun ctx => { ctx with useMT := true }) <| ematchTheorems thms
     withReader (fun ctx => { ctx with useMT := false }) <| ematchTheorems newThms
-  if (← checkMaxInstancesExceeded) then
+  if (← checkMaxInstancesExceeded <||> checkMaxEmatchExceeded) then
     return ()
   else
     go (← get).thms (← get).newThms |>.run'
@@ -335,6 +337,7 @@ def ematch : GoalM Unit := do
       thms         := s.thms ++ s.newThms
       newThms      := {}
       gmt          := s.gmt + 1
+      numEmatch    := s.numEmatch + 1
     }
 
 /-- Performs one round of E-matching, and assert new instances. -/

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

@@ -68,7 +68,7 @@ instance : Hashable CongrTheoremCacheKey where
   hash a := mixHash (unsafe ptrAddrUnsafe a.f).toUInt64 (hash a.numArgs)
 
 /-- State for the `GrindM` monad. -/
-structure CoreState where
+structure State where
   canon      : Canon.State := {}
   /-- `ShareCommon` (aka `Hashconsing`) state. -/
   scState    : ShareCommon.State.{0} ShareCommon.objectFactory := ShareCommon.State.mk _
@@ -88,7 +88,7 @@ private opaque MethodsRefPointed : NonemptyType.{0}
 private def MethodsRef : Type := MethodsRefPointed.type
 instance : Nonempty MethodsRef := MethodsRefPointed.property
 
-abbrev GrindM := ReaderT MethodsRef $ ReaderT Context $ StateRefT CoreState MetaM
+abbrev GrindM := ReaderT MethodsRef $ ReaderT Context $ StateRefT State MetaM
 
 /-- Returns the user-defined configuration options -/
 def getConfig : GrindM Grind.Config :=
@@ -207,7 +207,6 @@ structure ENode where
   generation : Nat := 0
   /-- Modification time -/
   mt : Nat := 0
-  -- TODO: see Lean 3 implementation
   deriving Inhabited, Repr
 
 def ENode.isCongrRoot (n : ENode) :=
@@ -375,6 +374,9 @@ structure Goal where
   thmMap       : EMatchTheorems
   /-- Number of theorem instances generated so far -/
   numInstances : Nat := 0
+  /-- Number of E-matching rounds performed in this goal so far. -/
+  -- Remark: it is always equal to `gmt` in the current implementation.
+  numEmatch    : Nat := 0
   /-- (pre-)instances found so far. It includes instances that failed to be instantiated. -/
   preInstances : PreInstanceSet := {}
   /-- new facts to be processed. -/
@@ -418,6 +420,10 @@ def addTheoremInstance (proof : Expr) (prop : Expr) (generation : Nat) : GoalM U
 def checkMaxInstancesExceeded : GoalM Bool := do
   return (← get).numInstances >= (← getConfig).instances
 
+/-- Returns `true` if the maximum number of E-matching rounds has been reached. -/
+def checkMaxEmatchExceeded : GoalM Bool := do
+  return (← get).numEmatch >= (← getConfig).ematch
+
 /--
 Returns `some n` if `e` has already been "internalized" into the
 Otherwise, returns `none`s.

tests/lean/run/grind_offset.lean

@@ -85,3 +85,70 @@ info: [grind.assert] foo (c + 1) = a
 example : foo (c + 1) = a → c = b + 1 → a = g (foo b) := by
   fail_if_success grind
   sorry
+
+set_option trace.grind.assert false
+
+/--
+info: [grind.ematch.instance] f.eq_2: f (x + 99).succ = g (f (x + 99))
+[grind.ematch.instance] f.eq_2: f (x + 98).succ = g (f (x + 98))
+-/
+#guard_msgs (info) in
+example : f (x + 100) = a → a = b := by
+  fail_if_success grind (ematch := 2)
+  sorry
+
+/--
+info: [grind.ematch.instance] f.eq_2: f (x + 99).succ = g (f (x + 99))
+[grind.ematch.instance] f.eq_2: f (x + 98).succ = g (f (x + 98))
+[grind.ematch.instance] f.eq_2: f (x + 97).succ = g (f (x + 97))
+[grind.ematch.instance] f.eq_2: f (x + 96).succ = g (f (x + 96))
+[grind.ematch.instance] f.eq_2: f (x + 95).succ = g (f (x + 95))
+-/
+#guard_msgs (info) in
+example : f (x + 100) = a → a = b := by
+  fail_if_success grind (ematch := 5)
+  sorry
+
+/--
+info: [grind.ematch.instance] f.eq_2: f (x + 99).succ = g (f (x + 99))
+[grind.ematch.instance] f.eq_2: f (x + 98).succ = g (f (x + 98))
+[grind.ematch.instance] f.eq_2: f (x + 97).succ = g (f (x + 97))
+[grind.ematch.instance] f.eq_2: f (x + 96).succ = g (f (x + 96))
+-/
+#guard_msgs (info) in
+example : f (x + 100) = a → a = b := by
+  fail_if_success grind (ematch := 100) (instances := 4)
+  sorry
+
+/--
+info: [grind.ematch.instance] f.eq_2: f (y + 9).succ = g (f (y + 9))
+[grind.ematch.instance] f.eq_2: f (x + 99).succ = g (f (x + 99))
+[grind.ematch.instance] f.eq_2: f (x + 98).succ = g (f (x + 98))
+[grind.ematch.instance] f.eq_2: f (y + 8).succ = g (f (y + 8))
+[grind.ematch.instance] f.eq_2: f (y + 7).succ = g (f (y + 7))
+-/
+#guard_msgs (info) in
+example : f (x + 100) = a → f (y + 10) = c → a = b := by
+  fail_if_success grind (ematch := 100) (instances := 5)
+  sorry
+
+/--
+info: [grind.ematch.instance] f.eq_2: f (x + 99).succ = g (f (x + 99))
+[grind.ematch.instance] f.eq_2: f (x + 98).succ = g (f (x + 98))
+-/
+#guard_msgs (info) in
+example : f (x + 100) = a → a = b := by
+  fail_if_success grind (gen := 2)
+  sorry
+
+/--
+info: [grind.ematch.instance] f.eq_2: f (x + 99).succ = g (f (x + 99))
+[grind.ematch.instance] f.eq_2: f (x + 98).succ = g (f (x + 98))
+[grind.ematch.instance] f.eq_2: f (x + 97).succ = g (f (x + 97))
+[grind.ematch.instance] f.eq_2: f (x + 96).succ = g (f (x + 96))
+[grind.ematch.instance] f.eq_2: f (x + 95).succ = g (f (x + 95))
+-/
+#guard_msgs (info) in
+example : f (x + 100) = a → a = b := by
+  fail_if_success grind (gen := 5)
+  sorry