fix: offset terms internalization (#6777)

This PR fixes a bug in the internalization of offset terms in the
`grind` tactic. For example, `grind` was failing to solve the following
example because of this bug.
```lean
example (f : Nat → Nat) : f (a + 1) = 1 → a = 0 → f 1 = 1 := by
  grind
```

src/Lean/Meta/Tactic/Grind.lean

@@ -72,6 +72,7 @@ builtin_initialize registerTraceClass `grind.debug.offset
 builtin_initialize registerTraceClass `grind.debug.offset.proof
 builtin_initialize registerTraceClass `grind.debug.ematch.pattern
 builtin_initialize registerTraceClass `grind.debug.beta
+builtin_initialize registerTraceClass `grind.debug.internalize
 builtin_initialize registerTraceClass `grind.debug.matchCond
 builtin_initialize registerTraceClass `grind.debug.matchCond.lambda
 

src/Lean/Meta/Tactic/Grind/Arith/Model.lean

@@ -38,8 +38,9 @@ def mkModel (goal : Goal) : MetaM (Array (Expr × Nat)) := do
     /-
     We should not include the assignment for auxiliary offset terms since
     they do not provide any additional information.
+    That said, the information is relevant for debugging `grind`.
     -/
-    if !(← isLitValue e) && (isNatOffset? e).isNone && isNatNum? e != some 0 then
+    if (!(← isLitValue e) && (isNatOffset? e).isNone && isNatNum? e != some 0) || grind.debug.get (← getOptions) then
       r := r.push (e, val)
   return r
 

src/Lean/Meta/Tactic/Grind/Arith/Offset.lean

@@ -270,7 +270,13 @@ private def isEqParent (parent? : Option Expr) : Bool := Id.run do
   let some parent := parent? | return false
   return parent.isEq
 
+private def alreadyInternalized (e : Expr) : GoalM Bool := do
+  let s ← get'
+  return s.cnstrs.contains { expr := e } || s.nodeMap.contains { expr := e }
+
 def internalize (e : Expr) (parent? : Option Expr) : GoalM Unit := do
+  if (← alreadyInternalized e) then
+    return ()
   let z ← getNatZeroExpr
   if let some c := isNatOffsetCnstr? e z then
     internalizeCnstr e c

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

@@ -114,7 +114,6 @@ private def preprocessGroundPattern (e : Expr) : GoalM Expr := do
 private def mkENode' (e : Expr) (generation : Nat) : GoalM Unit :=
   mkENodeCore e (ctor := false) (interpreted := false) (generation := generation)
 
-mutual
 /-- Internalizes the nested ground terms in the given pattern. -/
 private partial def internalizePattern (pattern : Expr) (generation : Nat) : GoalM Expr := do
   if pattern.isBVar || isPatternDontCare pattern then
@@ -127,7 +126,7 @@ private partial def internalizePattern (pattern : Expr) (generation : Nat) : Goa
     return mkAppN f (← args.mapM (internalizePattern · generation))
 
 /-- Internalizes the `MatchCond` gadget. -/
-private partial def internalizeMatchCond (matchCond : Expr) (generation : Nat) : GoalM Unit := do
+private def internalizeMatchCond (matchCond : Expr) (generation : Nat) : GoalM Unit := do
   mkENode' matchCond generation
   let (lhss, e') ← collectMatchCondLhssAndAbstract matchCond
   lhss.forM fun lhs => do internalize lhs generation; registerParent matchCond lhs
@@ -137,7 +136,7 @@ private partial def internalizeMatchCond (matchCond : Expr) (generation : Nat) :
   trace[grind.debug.matchCond.lambda] "auxiliary application{indentExpr e'}"
   pushEq matchCond e' (← mkEqRefl matchCond)
 
-partial def activateTheorem (thm : EMatchTheorem) (generation : Nat) : GoalM Unit := do
+def activateTheorem (thm : EMatchTheorem) (generation : Nat) : GoalM Unit := do
   -- Recall that we use the proof as part of the key for a set of instances found so far.
   -- We don't want to use structural equality when comparing keys.
   let proof ← shareCommon thm.proof
@@ -149,7 +148,7 @@ partial def activateTheorem (thm : EMatchTheorem) (generation : Nat) : GoalM Uni
 If `Config.matchEqs` is set to `true`, and `f` is `match`-auxiliary function,
 adds its equations to `newThms`.
 -/
-private partial def addMatchEqns (f : Expr) (generation : Nat) : GoalM Unit := do
+private def addMatchEqns (f : Expr) (generation : Nat) : GoalM Unit := do
   if !(← getConfig).matchEqs then return ()
   let .const declName _ := f | return ()
   if !(← isMatcher declName) then return ()
@@ -159,7 +158,7 @@ private partial def addMatchEqns (f : Expr) (generation : Nat) : GoalM Unit := d
     -- We disable pattern normalization to prevent the `match`-expression to be reduced.
     activateTheorem (← mkEMatchEqTheorem eqn (normalizePattern := false)) generation
 
-private partial def activateTheoremPatterns (fName : Name) (generation : Nat) : GoalM Unit := do
+private def activateTheoremPatterns (fName : Name) (generation : Nat) : GoalM Unit := do
   if let some (thms, thmMap) := (← get).thmMap.retrieve? fName then
     modify fun s => { s with thmMap }
     let appMap := (← get).appMap
@@ -173,8 +172,21 @@ private partial def activateTheoremPatterns (fName : Name) (generation : Nat) :
           trace_goal[grind.ematch] "reinsert `{thm.origin.key}`"
           modify fun s => { s with thmMap := s.thmMap.insert thm }
 
-partial def internalize (e : Expr) (generation : Nat) (parent? : Option Expr := none) : GoalM Unit := do
-  if (← alreadyInternalized e) then return ()
+
+@[export lean_grind_internalize]
+private partial def internalizeImpl (e : Expr) (generation : Nat) (parent? : Option Expr := none) : GoalM Unit := do
+  if (← alreadyInternalized e) then
+    trace_goal[grind.debug.internalize] "already internalized: {e}"
+    /-
+    Even if `e` has already been internalized, we must check whether it has also been internalized in
+    the satellite solvers. For example, suppose we have already internalized the term `f (a + 1)`.
+    The `1` in this term is treated as an offset for the offset term `a + 1` by the arithmetic module, and
+    only nodes for `a` and `a+1` are created. However, an ENode for `1` is created here.
+    Later, if we try to internalize `f 1`, the arithmetic module must create a node for `1`.
+    Otherwise, it will not be able to propagate that `a + 1 = 1` when `a = 0`
+    -/
+    Arith.internalize e parent?
+    return ()
   trace_goal[grind.internalize] "{e}"
   match e with
   | .bvar .. => unreachable!
@@ -183,10 +195,10 @@ partial def internalize (e : Expr) (generation : Nat) (parent? : Option Expr :=
   | .forallE _ d b _ =>
     mkENode' e generation
     if (← isProp d <&&> isProp e) then
-      internalize d generation e
+      internalizeImpl d generation e
       registerParent e d
       unless b.hasLooseBVars do
-        internalize b generation e
+        internalizeImpl b generation e
         registerParent e b
       propagateUp e
   | .lit .. | .const .. =>
@@ -215,17 +227,17 @@ partial def internalize (e : Expr) (generation : Nat) (parent? : Option Expr :=
         -- We only internalize the proposition. We can skip the proof because of
         -- proof irrelevance
         let c := args[0]!
-        internalize c generation e
+        internalizeImpl c generation e
         registerParent e c
       else if f.isConstOf ``ite && args.size == 5 then
         let c := args[1]!
-        internalize c generation e
+        internalizeImpl c generation e
         registerParent e c
       else
         if let .const fName _ := f then
           activateTheoremPatterns fName generation
         else
-          internalize f generation e
+          internalizeImpl f generation e
         registerParent e f
         for h : i in [: args.size] do
           let arg := args[i]
@@ -238,6 +250,4 @@ partial def internalize (e : Expr) (generation : Nat) (parent? : Option Expr :=
       propagateUp e
       propagateBetaForNewApp e
 
-end
-
 end Lean.Meta.Grind

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

@@ -740,6 +740,10 @@ It assumes `a` and `b` are in the same equivalence class.
 @[extern "lean_grind_mk_heq_proof"]
 opaque mkHEqProof (a b : Expr) : GoalM Expr
 
+-- Forward definition
+@[extern "lean_grind_internalize"]
+opaque internalize (e : Expr) (generation : Nat) (parent? : Option Expr := none) : GoalM Unit
+
 /--
 Returns a proof that `a = b` if they have the same type. Otherwise, returns a proof of `HEq a b`.
 It assumes `a` and `b` are in the same equivalence class.

tests/lean/run/grind_pre.lean

@@ -97,6 +97,7 @@ x✝ : ¬g (i + 1) j ⋯ = i + j + 1
     [assign] i + j := 1
 -/
 #guard_msgs (error) in
+set_option grind.debug false in
 example (i j : Nat) (h : i + 1 > j + 1) : g (i+1) j = f ((fun x => x) i) + f j + 1 := by
   grind
 
@@ -192,10 +193,8 @@ info: [grind.issues] found congruence between
     and
       f a
     but functions have different types
----
-warning: declaration uses 'sorry'
 -/
-#guard_msgs in
+#guard_msgs (info) in
 set_option trace.grind.issues true in
 set_option trace.grind.debug.proof false in
 example (f : Nat → Bool) (g : Int → Bool) (a : Nat) (b : Int) : HEq f g → HEq a b → f a = g b := by

tests/lean/run/grind_t1.lean

@@ -357,3 +357,6 @@ set_option trace.grind.issues true in
 example : (if n + 2 < m then a else b) = (if n + 1 < m then c else d) := by
   fail_if_success grind (splits := 0)
   sorry
+
+example (f : Nat → Nat) : f (a + 1) = 1 → a = 0 → f 1 = 1 := by
+  grind