chore: fix test

refactor: remove unnecessary `Cnstr.le` field, minimize number of `markAsOffsetTerm`

fix: do not inclued assignment for auxiliary offset terms in the model

fix: ensure e-graph is in a consistent state before propagating constraints

feat: offset equality support

feat: equality propagation from offset constraint module to core

chore: fix tests

src/Init/Grind/Offset.lean

@@ -80,4 +80,13 @@ theorem Nat.ro_eq_false_of_lo (u v k₁ k₂ : Nat) : isLt k₂ k₁ = true →
 theorem Nat.lo_eq_false_of_ro (u v k₁ k₂ : Nat) : isLt k₁ k₂ = true → u ≤ v + k₁ → (v + k₂ ≤ u) = False := by
   simp [isLt]; omega
 
+/-!
+Helper theorems for equality propagation
+-/
+
+theorem Nat.le_of_eq_1 (u v : Nat) : u = v → u ≤ v := by omega
+theorem Nat.le_of_eq_2 (u v : Nat) : u = v → v ≤ u := by omega
+theorem Nat.eq_of_le_of_le (u v : Nat) : u ≤ v → v ≤ u → u = v := by omega
+theorem Nat.le_offset (a k : Nat) : k ≤ a + k := by omega
+
 end Lean.Grind

src/Lean/Meta/Tactic/Grind.lean

@@ -48,6 +48,9 @@ builtin_initialize registerTraceClass `grind.offset.dist
 builtin_initialize registerTraceClass `grind.offset.internalize
 builtin_initialize registerTraceClass `grind.offset.internalize.term (inherited := true)
 builtin_initialize registerTraceClass `grind.offset.propagate
+builtin_initialize registerTraceClass `grind.offset.eq
+builtin_initialize registerTraceClass `grind.offset.eq.to (inherited := true)
+builtin_initialize registerTraceClass `grind.offset.eq.from (inherited := true)
 
 /-! Trace options for `grind` developers -/
 builtin_initialize registerTraceClass `grind.debug

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

@@ -8,7 +8,7 @@ import Lean.Meta.Tactic.Grind.Arith.Offset
 
 namespace Lean.Meta.Grind.Arith
 
-def internalize (e : Expr) : GoalM Unit := do
-  Offset.internalizeCnstr e
+def internalize (e : Expr) (parent? : Option Expr) : GoalM Unit := do
+  Offset.internalize e parent?
 
 end Lean.Meta.Grind.Arith

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

@@ -6,6 +6,7 @@ Authors: Leonardo de Moura
 prelude
 import Lean.Meta.Basic
 import Lean.Meta.Tactic.Grind.Types
+import Lean.Meta.Tactic.Grind.Util
 
 namespace Lean.Meta.Grind.Arith.Offset
 /-- Construct a model that statisfies all offset constraints -/
@@ -33,7 +34,13 @@ def mkModel (goal : Goal) : MetaM (Array (Expr × Nat)) := do
   for u in [:nodes.size] do
     let some val := pre[u]! | unreachable!
     let val := (val - min).toNat
-    r := r.push (nodes[u]!, val)
+    let e := nodes[u]!
+    /-
+    We should not include the assignment for auxiliary offset terms since
+    they do not provide any additional information.
+    -/
+    if (isNatOffset? e).isNone && isNatNum? e != some 0 then
+      r := r.push (e, val)
   return r
 
 end Lean.Meta.Grind.Arith.Offset

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

@@ -48,6 +48,7 @@ def mkNode (expr : Expr) : GoalM NodeId := do
     targets := s.targets.push {}
     proofs  := s.proofs.push {}
   }
+  markAsOffsetTerm expr
   return nodeId
 
 private def getExpr (u : NodeId) : GoalM Expr := do
@@ -59,6 +60,11 @@ private def getDist? (u v : NodeId) : GoalM (Option Int) := do
 private def getProof? (u v : NodeId) : GoalM (Option ProofInfo) := do
   return (← get').proofs[u]!.find? v
 
+private def getNodeId (e : Expr) : GoalM NodeId := do
+  let some nodeId := (← get').nodeMap.find? { expr := e }
+    | throwError "internal `grind` error, term has not been internalized by offset module{indentExpr e}"
+  return nodeId
+
 /--
 Returns a proof for `u + k ≤ v` (or `u ≤ v + k`) where `k` is the
 shortest path between `u` and `v`.
@@ -160,10 +166,24 @@ private def updateCnstrsOf (u v : NodeId) (f : Cnstr NodeId → Expr → GoalM B
         return !(← f c e)
     modify' fun s => { s with cnstrsOf := s.cnstrsOf.insert (u, v) cs' }
 
+/-- Equality propagation. -/
+private def propagateEq (u v : NodeId) (k : Int) : GoalM Unit := do
+  if k != 0 then return ()
+  let some k' ← getDist? v u | return ()
+  if k' != 0 then return ()
+  let ue ← getExpr u
+  let ve ← getExpr v
+  if (← isEqv ue ve) then return ()
+  let huv ← mkProofForPath u v
+  let hvu ← mkProofForPath v u
+  trace[grind.offset.eq.from] "{ue}, {ve}"
+  pushEq ue ve <| mkApp4 (mkConst ``Grind.Nat.eq_of_le_of_le) ue ve huv hvu
+
 /-- Performs constraint propagation. -/
 private def propagateAll (u v : NodeId) (k : Int) : GoalM Unit := do
   updateCnstrsOf u v fun c e => return !(← propagateTrue u v k c e)
   updateCnstrsOf v u fun c e => return !(← propagateFalse u v k c e)
+  propagateEq u v k
 
 /--
 If `isShorter u v k`, updates the shortest distance between `u` and `v`.
@@ -203,8 +223,7 @@ where
         /- Check whether new path: `i -(k₁)-> u -(k)-> v -(k₂) -> j` is shorter -/
         updateIfShorter i j (k₁+k+k₂) v
 
-def internalizeCnstr (e : Expr) : GoalM Unit := do
-  let some c := isNatOffsetCnstr? e | return ()
+private def internalizeCnstr (e : Expr) (c : Cnstr Expr) : GoalM Unit := do
   let u ← mkNode c.u
   let v ← mkNode c.v
   let c := { c with u, v }
@@ -222,6 +241,70 @@ def internalizeCnstr (e : Expr) : GoalM Unit := do
       s.cnstrsOf.insert (u, v) cs
   }
 
+private def getZeroNode : GoalM NodeId := do
+  mkNode (← getNatZeroExpr)
+
+/-- Internalize `e` of the form `b + k` -/
+private def internalizeTerm (e : Expr) (b : Expr) (k : Nat) : GoalM Unit := do
+  -- `e` is of the form `b + k`
+  let u ← mkNode e
+  let v ← mkNode b
+  -- `u = v + k`. So, we add edges for `u ≤ v + k` and `v + k ≤ u`.
+  let h := mkApp (mkConst ``Nat.le_refl) e
+  addEdge u v k h
+  addEdge v u (-k) h
+  -- `0 + k ≤ u`
+  let z ← getZeroNode
+  addEdge z u (-k) <| mkApp2 (mkConst ``Grind.Nat.le_offset) b (toExpr k)
+
+/--
+Returns `true`, if `parent?` is relevant for internalization.
+For example, we do not want to internalize an offset term that
+is the child of an addition. This kind of term will be processed by the
+more general linear arithmetic module.
+-/
+private def isRelevantParent (parent? : Option Expr) : GoalM Bool := do
+  let some parent := parent? | return false
+  let z ← getNatZeroExpr
+  return !isNatAdd parent && (isNatOffsetCnstr? parent z).isNone
+
+private def isEqParent (parent? : Option Expr) : Bool := Id.run do
+  let some parent := parent? | return false
+  return parent.isEq
+
+def internalize (e : Expr) (parent? : Option Expr) : GoalM Unit := do
+  let z ← getNatZeroExpr
+  if let some c := isNatOffsetCnstr? e z then
+    internalizeCnstr e c
+  else if (← isRelevantParent parent?) then
+    if let some (b, k) := isNatOffset? e then
+      internalizeTerm e b k
+    else if let some k := isNatNum? e then
+      -- core module has support for detecting equality between literals
+      unless isEqParent parent? do
+        internalizeTerm e z k
+
+@[export lean_process_new_offset_eq]
+def processNewOffsetEqImpl (a b : Expr) : GoalM Unit := do
+  unless isSameExpr a b do
+    trace[grind.offset.eq.to] "{a}, {b}"
+    let u ← getNodeId a
+    let v ← getNodeId b
+    let h ← mkEqProof a b
+    addEdge u v 0 <| mkApp3 (mkConst ``Grind.Nat.le_of_eq_1) a b h
+    addEdge v u 0 <| mkApp3 (mkConst ``Grind.Nat.le_of_eq_2) a b h
+
+@[export lean_process_new_offset_eq_lit]
+def processNewOffsetEqLitImpl (a b : Expr) : GoalM Unit := do
+  unless isSameExpr a b do
+    trace[grind.offset.eq.to] "{a}, {b}"
+    let some k := isNatNum? b | unreachable!
+    let u ← getNodeId a
+    let z ← mkNode (← getNatZeroExpr)
+    let h ← mkEqProof a b
+    addEdge u z k <| mkApp3 (mkConst ``Grind.Nat.le_of_eq_1) a b h
+    addEdge z u (-k) <| mkApp3 (mkConst ``Grind.Nat.le_of_eq_2) a b h
+
 def traceDists : GoalM Unit := do
   let s ← get'
   for u in [:s.targets.size], es in s.targets.toArray do
@@ -231,13 +314,12 @@ def traceDists : GoalM Unit := do
 def Cnstr.toExpr (c : Cnstr NodeId) : GoalM Expr := do
   let u := (← get').nodes[c.u]!
   let v := (← get').nodes[c.v]!
-  let mk := if c.le then mkNatLE else mkNatEq
   if c.k == 0 then
-    return mk u v
+    return mkNatLE u v
   else if c.k < 0 then
-    return mk (mkNatAdd u (Lean.toExpr ((-c.k).toNat))) v
+    return mkNatLE (mkNatAdd u (Lean.toExpr ((-c.k).toNat))) v
   else
-    return mk u (mkNatAdd v (Lean.toExpr c.k.toNat))
+    return mkNatLE u (mkNatAdd v (Lean.toExpr c.k.toNat))
 
 def checkInvariants : GoalM Unit := do
   let s ← get'

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

@@ -82,7 +82,7 @@ def mkOfNegEqFalse (nodes : PArray Expr) (c : Cnstr NodeId) (h : Expr) : Expr :=
   let v := nodes[c.v]!
   if c.k == 0 then
     mkApp3 (mkConst ``Nat.of_le_eq_false) u v h
-  else if c.k == -1 && c.le then
+  else if c.k == -1 then
     mkApp3 (mkConst ``Nat.of_lo_eq_false_1) u v h
   else if c.k < 0 then
     mkApp4 (mkConst ``Nat.of_lo_eq_false) u v (toExprN (-c.k)) h

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

@@ -32,6 +32,16 @@ def isNatAdd? (e : Expr) : Option (Expr × Expr) :=
   let_expr HAdd.hAdd _ _ _ i a b := e | none
   if isInstAddNat i then some (a, b) else none
 
+/--
+Returns `true` if `e` is of the form
+```
+@HAdd.hAdd Nat Nat Nat (instHAdd Nat instAddNat) _ _
+```
+-/
+def isNatAdd (e : Expr) : Bool :=
+  let_expr HAdd.hAdd _ _ _ i _ _ := e | false
+  isInstAddNat i
+
 /-- Returns `some k` if `e` `@OfNat.ofNat Nat _ (instOfNatNat k)` -/
 def isNatNum? (e : Expr) : Option Nat := Id.run do
   let_expr OfNat.ofNat _ _ inst := e | none
@@ -50,40 +60,43 @@ structure Offset.Cnstr (α : Type) where
   u  : α
   v  : α
   k  : Int := 0
-  le : Bool := true
   deriving Inhabited
 
 def Offset.Cnstr.neg : Cnstr α → Cnstr α
-  | { u, v, k, le } => { u := v, v := u, le, k := -k - 1 }
+  | { u, v, k } => { u := v, v := u, k := -k - 1 }
 
 example (c : Offset.Cnstr α) : c.neg.neg = c := by
   cases c; simp [Offset.Cnstr.neg]; omega
 
 def Offset.toMessageData [inst : ToMessageData α] (c : Offset.Cnstr α) : MessageData :=
-  match c.k, c.le with
-  | .ofNat 0,   true  => m!"{c.u} ≤ {c.v}"
-  | .ofNat 0,   false => m!"{c.u} = {c.v}"
-  | .ofNat k,   true  => m!"{c.u} ≤ {c.v} + {k}"
-  | .ofNat k,   false => m!"{c.u} = {c.v} + {k}"
-  | .negSucc k, true  => m!"{c.u} + {k + 1} ≤ {c.v}"
-  | .negSucc k, false => m!"{c.u} + {k + 1} = {c.v}"
+  match c.k with
+  | .ofNat 0   => m!"{c.u} ≤ {c.v}"
+  | .ofNat k   => m!"{c.u} ≤ {c.v} + {k}"
+  | .negSucc k => m!"{c.u} + {k + 1} ≤ {c.v}"
 
 instance : ToMessageData (Offset.Cnstr Expr) where
   toMessageData c := Offset.toMessageData c
 
-/-- Returns `some cnstr` if `e` is offset constraint. -/
-def isNatOffsetCnstr? (e : Expr) : Option (Offset.Cnstr Expr) :=
+/--
+Returns `some cnstr` if `e` is offset constraint.
+Remark: `z` is `0` numeral. It is an extra argument because we
+want to be able to provide the one that has already been internalized.
+-/
+def isNatOffsetCnstr? (e : Expr) (z : Expr) : Option (Offset.Cnstr Expr) :=
   match_expr e with
-  | LE.le _ inst a b => if isInstLENat inst then go a b true else none
-  | Eq α a b => if isNatType α then go a b false else none
+  | LE.le _ inst a b => if isInstLENat inst then go a b else none
   | _ => none
 where
-  go (u v : Expr) (le : Bool) :=
+  go (u v : Expr) :=
     if let some (u, k) := isNatOffset? u then
-      some { u, k := - k, v, le }
+      some { u, k := - k, v }
     else if let some (v, k) := isNatOffset? v then
-      some { u, v, k := k, le }
+      some { u, v, k }
+    else if let some k := isNatNum? u then
+      some { u := z, v, k := - k }
+    else if let some k := isNatNum? v then
+      some { u, v := z, k }
     else
-      some { u, v, le }
+      some { u, v }
 
 end Lean.Meta.Grind.Arith

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

@@ -10,6 +10,7 @@ import Lean.Meta.Tactic.Grind.Types
 import Lean.Meta.Tactic.Grind.Inv
 import Lean.Meta.Tactic.Grind.PP
 import Lean.Meta.Tactic.Grind.Ctor
+import Lean.Meta.Tactic.Grind.Util
 import Lean.Meta.Tactic.Grind.Internalize
 
 namespace Lean.Meta.Grind
@@ -86,6 +87,26 @@ private partial def updateMT (root : Expr) : GoalM Unit := do
       setENode parent { node with mt := gmt }
       updateMT parent
 
+/--
+Helper function for combining `ENode.offset?` fields and propagating an equality
+to the offset constraint module.
+-/
+private def propagateOffsetEq (rhsRoot lhsRoot : ENode) : GoalM Unit := do
+  match lhsRoot.offset? with
+  | some lhsOffset =>
+    if let some rhsOffset := rhsRoot.offset? then
+      Arith.processNewOffsetEq lhsOffset rhsOffset
+    else if isNatNum rhsRoot.self then
+      Arith.processNewOffsetEqLit lhsOffset rhsRoot.self
+    else
+      -- We have to retrieve the node because other fields have been updated
+      let rhsRoot ← getENode rhsRoot.self
+      setENode rhsRoot.self { rhsRoot with offset? := lhsOffset }
+  | none =>
+    if isNatNum lhsRoot.self then
+    if let some rhsOffset := rhsRoot.offset? then
+      Arith.processNewOffsetEqLit rhsOffset lhsRoot.self
+
 private partial def addEqStep (lhs rhs proof : Expr) (isHEq : Bool) : GoalM Unit := do
   let lhsNode ← getENode lhs
   let rhsNode ← getENode rhs
@@ -146,17 +167,18 @@ where
       next := rhsRoot.next
     }
     setENode rhsNode.root { rhsRoot with
-      next := lhsRoot.next
-      size := rhsRoot.size + lhsRoot.size
+      next       := lhsRoot.next
+      size       := rhsRoot.size + lhsRoot.size
       hasLambdas := rhsRoot.hasLambdas || lhsRoot.hasLambdas
       heqProofs  := isHEq || rhsRoot.heqProofs || lhsRoot.heqProofs
     }
     copyParentsTo parents rhsNode.root
+    unless (← isInconsistent) do
+      updateMT rhsRoot.self
+    propagateOffsetEq rhsRoot lhsRoot
     unless (← isInconsistent) do
       for parent in parents do
         propagateUp parent
-    unless (← isInconsistent) do
-      updateMT rhsRoot.self
 
   updateRoots (lhs : Expr) (rootNew : Expr) : GoalM Unit := do
     traverseEqc lhs fun n =>
@@ -205,9 +227,10 @@ private def storeFact (fact : Expr) : GoalM Unit := do
 
 /-- Internalizes `lhs` and `rhs`, and then adds equality `lhs = rhs`. -/
 def addNewEq (lhs rhs proof : Expr) (generation : Nat) : GoalM Unit := do
-  storeFact (← mkEq lhs rhs)
-  internalize lhs generation
-  internalize rhs generation
+  let eq ← mkEq lhs rhs
+  storeFact eq
+  internalize lhs generation eq
+  internalize rhs generation eq
   addEq lhs rhs proof
 
 /-- Adds a new `fact` justified by the given proof and using the given generation. -/
@@ -244,8 +267,8 @@ where
       internalize p generation
       addEq p (← getFalseExpr) (← mkEqFalse proof)
     else
-      internalize lhs generation
-      internalize rhs generation
+      internalize lhs generation p
+      internalize rhs generation p
       addEqCore lhs rhs proof isHEq
 
 /-- Adds a new hypothesis. -/

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

@@ -105,7 +105,7 @@ private partial def internalizePattern (pattern : Expr) (generation : Nat) : Goa
     return pattern
   else if let some e := groundPattern? pattern then
     let e ← shareCommon (← canon (← normalizeLevels (← unfoldReducible e)))
-    internalize e generation
+    internalize e generation none
     return mkGroundPattern e
   else pattern.withApp fun f args => do
     return mkAppN f (← args.mapM (internalizePattern · generation))
@@ -146,7 +146,7 @@ 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) : GoalM Unit := do
+partial def internalize (e : Expr) (generation : Nat) (parent? : Option Expr := none) : GoalM Unit := do
   if (← alreadyInternalized e) then return ()
   trace_goal[grind.internalize] "{e}"
   match e with
@@ -157,10 +157,10 @@ partial def internalize (e : Expr) (generation : Nat) : GoalM Unit := do
   | .forallE _ d b _ =>
     mkENodeCore e (ctor := false) (interpreted := false) (generation := generation)
     if (← isProp d <&&> isProp e) then
-      internalize d generation
+      internalize d generation e
       registerParent e d
       unless b.hasLooseBVars do
-        internalize b generation
+        internalize b generation e
         registerParent e b
       propagateUp e
   | .lit .. | .const .. =>
@@ -174,6 +174,7 @@ partial def internalize (e : Expr) (generation : Nat) : GoalM Unit := do
     if (← isLitValue e) then
       -- We do not want to internalize the components of a literal value.
       mkENode e generation
+      Arith.internalize e parent?
     else e.withApp fun f args => do
       checkAndAddSplitCandidate e
       pushCastHEqs e
@@ -182,22 +183,22 @@ partial def internalize (e : Expr) (generation : Nat) : GoalM Unit := do
         -- We only internalize the proposition. We can skip the proof because of
         -- proof irrelevance
         let c := args[0]!
-        internalize c generation
+        internalize c generation e
         registerParent e c
       else
         if let .const fName _ := f then
           activateTheoremPatterns fName generation
         else
-          internalize f generation
+          internalize f generation e
           registerParent e f
         for h : i in [: args.size] do
           let arg := args[i]
-          internalize arg generation
+          internalize arg generation e
           registerParent e arg
       mkENode e generation
       addCongrTable e
       updateAppMap e
-      Arith.internalize e
+      Arith.internalize e parent?
       propagateUp e
 end
 

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

@@ -40,17 +40,20 @@ def GrindM.run (x : GrindM α) (mainDeclName : Name) (config : Grind.Config) (fa
   let scState := ShareCommon.State.mk _
   let (falseExpr, scState) := ShareCommon.State.shareCommon scState (mkConst ``False)
   let (trueExpr, scState)  := ShareCommon.State.shareCommon scState (mkConst ``True)
+  let (natZExpr, scState)  := ShareCommon.State.shareCommon scState (mkNatLit 0)
   let simprocs ← Grind.getSimprocs
   let simp ← Grind.getSimpContext
-  x (← mkMethods fallback).toMethodsRef { mainDeclName, config, simprocs, simp } |>.run' { scState, trueExpr, falseExpr }
+  x (← mkMethods fallback).toMethodsRef { mainDeclName, config, simprocs, simp } |>.run' { scState, trueExpr, falseExpr, natZExpr }
 
 private def mkGoal (mvarId : MVarId) : GrindM Goal := do
   let trueExpr ← getTrueExpr
   let falseExpr ← getFalseExpr
+  let natZeroExpr ← getNatZeroExpr
   let thmMap ← getEMatchTheorems
   GoalM.run' { mvarId, thmMap } do
     mkENodeCore falseExpr (interpreted := true) (ctor := false) (generation := 0)
     mkENodeCore trueExpr (interpreted := true) (ctor := false) (generation := 0)
+    mkENodeCore natZeroExpr (interpreted := true) (ctor := false) (generation := 0)
 
 private def initCore (mvarId : MVarId) : GrindM (List Goal) := do
   mvarId.ensureProp

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

@@ -80,6 +80,7 @@ structure State where
   simpStats  : Simp.Stats := {}
   trueExpr   : Expr
   falseExpr  : Expr
+  natZExpr   : Expr
   /--
   Used to generate trace messages of the for `[grind] working on <tag>`,
   and implement the macro `trace_goal`.
@@ -104,6 +105,10 @@ def getTrueExpr : GrindM Expr := do
 def getFalseExpr : GrindM Expr := do
   return (← get).falseExpr
 
+/-- Returns the internalized `0 : Nat` numeral.  -/
+def getNatZeroExpr : GrindM Expr := do
+  return (← get).natZExpr
+
 def getMainDeclName : GrindM Name :=
   return (← readThe Context).mainDeclName
 
@@ -128,9 +133,9 @@ Applies hash-consing to `e`. Recall that all expressions in a `grind` goal have
 been hash-consed. We perform this step before we internalize expressions.
 -/
 def shareCommon (e : Expr) : GrindM Expr := do
-  modifyGet fun { canon, scState, nextThmIdx, congrThms, trueExpr, falseExpr, simpStats, lastTag } =>
+  modifyGet fun { canon, scState, nextThmIdx, congrThms, trueExpr, falseExpr, natZExpr, simpStats, lastTag } =>
     let (e, scState) := ShareCommon.State.shareCommon scState e
-    (e, { canon, scState, nextThmIdx, congrThms, trueExpr, falseExpr, simpStats, lastTag })
+    (e, { canon, scState, nextThmIdx, congrThms, trueExpr, falseExpr, natZExpr, simpStats, lastTag })
 
 /--
 Canonicalizes nested types, type formers, and instances in `e`.
@@ -202,13 +207,19 @@ structure ENode where
   on heterogeneous equality.
   -/
   heqProofs : Bool := false
-  /--
-  Unique index used for pretty printing and debugging purposes.
-  -/
+  /-- Unique index used for pretty printing and debugging purposes. -/
   idx : Nat := 0
+  /-- The generation in which this enode was created. -/
   generation : Nat := 0
   /-- Modification time -/
   mt : Nat := 0
+  /--
+  The `offset?` field is used to propagate equalities from the `grind` congruence closure module
+  to the offset constraints module. When `grind` merges two equivalence classes, and both have
+  an associated `offset?` set to `some e`, the equality is propagated. This field is
+  assigned during the internalization of offset terms.
+  -/
+  offset? : Option Expr := none
   deriving Inhabited, Repr
 
 def ENode.isCongrRoot (n : ENode) :=
@@ -643,6 +654,41 @@ def mkENode (e : Expr) (generation : Nat) : GoalM Unit := do
   let interpreted ← isInterpreted e
   mkENodeCore e interpreted ctor generation
 
+/--
+Notify the offset constraint module that `a = b` where
+`a` and `b` are terms that have been internalized by this module.
+-/
+@[extern "lean_process_new_offset_eq"] -- forward definition
+opaque Arith.processNewOffsetEq (a b : Expr) : GoalM Unit
+
+/--
+Notify the offset constraint module that `a = k` where
+`a` is term that has been internalized by this module,
+and `k` is a numeral.
+-/
+@[extern "lean_process_new_offset_eq_lit"] -- forward definition
+opaque Arith.processNewOffsetEqLit (a k : Expr) : GoalM Unit
+
+/-- Returns `true` if `e` is a numeral and has type `Nat`. -/
+def isNatNum (e : Expr) : Bool := Id.run do
+  let_expr OfNat.ofNat _ _ inst := e | false
+  let_expr instOfNatNat _ := inst | false
+  true
+
+/--
+Marks `e` as a term of interest to the offset constraint module.
+If the root of `e`s equivalence class has already a term of interest,
+a new equality is propagated to the offset module.
+-/
+def markAsOffsetTerm (e : Expr) : GoalM Unit := do
+  let root ← getRootENode e
+  if let some e' := root.offset? then
+    Arith.processNewOffsetEq e e'
+  else if isNatNum root.self && !isSameExpr e root.self then
+    Arith.processNewOffsetEqLit e root.self
+  else
+    setENode root.self { root with offset? := some e }
+
 /-- Returns `true` is `e` is the root of its congruence class. -/
 def isCongrRoot (e : Expr) : GoalM Bool := do
   return (← getENode e).isCongrRoot

tests/lean/run/grind_offset.lean

@@ -83,8 +83,7 @@ info: [grind.assert] foo (c + 1) = a
 -/
 #guard_msgs (info) in
 example : foo (c + 1) = a → c = b + 1 → a = g (foo b) := by
-  fail_if_success grind
-  sorry
+  grind
 
 set_option trace.grind.assert false
 

tests/lean/run/grind_offset_cnstr.lean

@@ -352,3 +352,53 @@ example (p r : Prop) (a b : Nat) : (c + 1 ≤ a ↔ p) → (c + 2 ≤ a + 1 ↔
 set_option trace.grind.split true in
 example (p r : Prop) (a b : Nat) : (c + 5 ≤ a ↔ p) → (c + 4 ≤ a ↔ r) → a ≤ b → b ≤ c + 3 → ¬p ∧ ¬r := by
   grind (splits := 0)
+
+example (a b c d: Nat) : a ≤ b → b + 2 = c → c < d → a + 2 < d := by
+  grind
+
+example (a b c : Nat) : a + 2 = b → b + 3 = c → a + 5 ≤ c := by
+  grind
+
+example (a b c : Nat) : a + 2 = b → c ≤ a + 2 → a + 2 ≤ c → c = b := by
+  grind
+
+example (a b c : Nat) : a + 2 = b → b + 3 = c → a + 5 = c := by
+  grind
+
+example (f : Nat → Nat) (a b c d e : Nat) :
+        f (a + 3) = b →
+        f (c + 1) = d →
+        c ≤ a + 2 →
+        a + 1 ≤ e →
+        e < c →
+        b = d := by
+  grind
+
+example (a : Nat) : a < 2 → a < 5 := by
+  grind
+
+example (a b : Nat) : 2 < a → a ≤ b → 2 < b := by
+  grind
+
+example (a b : Nat) : 2 < a → a ≤ b → 0 < b := by
+  grind
+
+example (f : Nat → Nat) : f 1 = a → b ≤ 1 → b ≥ 1 → f b = a := by
+  grind
+
+example (f : Nat → Nat) : f 2 = a → b ≤ 1 → b ≥ 1 → c = b + 1 → f c = a := by
+  grind
+
+example (a : Nat) : a < 2 → a = 5 → False := by
+  grind
+
+example (a : Nat) : a < 2 → a = b → b = c → c = 5 → False := by
+  grind
+
+#guard_msgs (info) in -- none of the numerals should be internalized by the offset module
+set_option trace.grind.offset.internalize true in
+example (a b c d e : Nat) : a = 1 → b = 2 → c = 3 → d = 4 → e = 5 → a ≠ e := by
+  grind
+
+example (a b : Nat) : a + 1 = b → b = 0 → False := by
+  grind

tests/lean/run/grind_pre.lean

@@ -77,8 +77,7 @@ x✝ : ¬g (i + 1) j ⋯ = i + j + 1
   [prop] g (i + 1) j ⋯ = i + j + 1[offset] Assignment satisfying offset contraints
   [assign] j := 0
   [assign] i := 1
-  [assign] g (i + 1) j ⋯ := 0
-  [assign] i + j := 0
+  [assign] i + j := 1
 -/
 #guard_msgs (error) in
 example (i j : Nat) (h : i + 1 > j + 1) : g (i+1) j = f ((fun x => x) i) + f j + 1 := by