feat: equality propagation in grind order (#11047)

This PR implements (nested term) equality propagation in `grind order`.
That is, it propagates implied equalities from `grind order` back to the
`grind` core. Examples:
```lean
open Lean Grind Std

example [LE α] [IsPartialOrder α] (a b : α) (f : α → Nat) : a ≤ b → b ≤ c → c ≤ a → f a = f b := by
  grind (splits := 0)

example [CommRing α] [LE α] [LT α] [LawfulOrderLT α] [IsPartialOrder α] [OrderedRing α]
    (a b : α) (f : α → Int) : a ≤ b + 1 → b ≤ a - 1 → f a = f (2 + b - 1) := by
  grind -mbtc -lia -linarith (splits := 0)

example (a b : Int) (f : Int → Int) : a ≤ b + 1 → b ≤ a - 1 → f a = f (2 + b - 1) := by
  grind -mbtc -lia -linarith (splits := 0)
```

Author: leodemoura

src/Init/Grind/Order.lean

@@ -55,6 +55,16 @@ theorem le_of_eq_2_k {α} [LE α] [LT α] [Std.LawfulOrderLT α] [Std.IsPreorder
   rw [Ring.intCast_zero, Semiring.add_zero]
   intro h; subst a; apply Std.IsPreorder.le_refl
 
+theorem le_of_offset_eq_1_k {α} [LE α] [LT α] [Std.LawfulOrderLT α] [Std.IsPreorder α] [Ring α] [OrderedRing α]
+    {a b : α} {k : Int} : a = b + k → a ≤ b + k := by
+  intro h; subst a; apply Std.IsPreorder.le_refl
+
+theorem le_of_offset_eq_2_k {α} [LE α] [LT α] [Std.LawfulOrderLT α] [Std.IsPreorder α] [Ring α] [OrderedRing α]
+    {a b : α} {k : Int} : a = b + k → b ≤ a + (-k : Int) := by
+  intro h; subst a
+  rw [Ring.intCast_neg, Semiring.add_assoc, Semiring.add_comm (α := α) k, Ring.neg_add_cancel, Semiring.add_zero]
+  apply Std.IsPreorder.le_refl
+
 theorem le_of_not_le {α} [LE α] [Std.IsLinearPreorder α]
     {a b : α} : ¬ a ≤ b → b ≤ a := by
   intro h

src/Lean/Meta/Tactic/Grind/Arith/CommRing/Poly.lean

@@ -200,6 +200,10 @@ def Poly.isZero : Poly → Bool
   | .num 0 => true
   | _ => false
 
+def Poly.getConst : Poly → Int
+  | .num k => k
+  | .add _ _ p => p.getConst
+
 def Poly.checkCoeffs : Poly → Bool
   | .num _ => true
   | .add k _ p => k != 0 && checkCoeffs p

src/Lean/Meta/Tactic/Grind/Arith/CommRing/Proof.lean

@@ -413,6 +413,16 @@ def mkEqIffProof (lhs rhs lhs' rhs' : RingExpr) : RingM Expr := do
   let h := mkApp2 (mkConst ``Grind.CommRing.eq_norm_expr [ring.u]) ring.type ring.commRingInst
   return mkApp6 h ctx (toExpr lhs) (toExpr rhs) (toExpr lhs') (toExpr rhs') eagerReflBoolTrue
 
+/--
+Given `e` and `e'` s.t. `e.toPoly == e'.toPoly`, returns a proof that `e.denote ctx = e'.denote ctx`
+-/
+def mkTermEqProof (e e' : RingExpr) : RingM Expr := do
+  let ring ← getCommRing
+  let { lhs, lhs', vars, .. } := norm ring.vars e (.num 0) e' (.num 0)
+  let ctx ← toContextExpr vars
+  let h := mkApp2 (mkConst ``Grind.CommRing.Expr.eq_of_toPoly_eq [ring.u]) ring.type ring.commRingInst
+  return mkApp4 h ctx (toExpr lhs) (toExpr lhs') eagerReflBoolTrue
+
 def mkNonCommLeIffProof (leInst ltInst isPreorderInst orderedRingInst : Expr) (lhs rhs lhs' rhs' : RingExpr) : NonCommRingM Expr := do
   let ring ← getRing
   let { lhs, rhs, lhs', rhs', vars } := norm ring.vars lhs rhs lhs' rhs'
@@ -434,4 +444,14 @@ def mkNonCommEqIffProof (lhs rhs lhs' rhs' : RingExpr) : NonCommRingM Expr := do
   let h := mkApp2 (mkConst ``Grind.CommRing.eq_norm_expr_nc [ring.u]) ring.type ring.ringInst
   return mkApp6 h ctx (toExpr lhs) (toExpr rhs) (toExpr lhs') (toExpr rhs') eagerReflBoolTrue
 
+/--
+Given `e` and `e'` s.t. `e.toPoly_nc == e'.toPoly_nc`, returns a proof that `e.denote ctx = e'.denote ctx`
+-/
+def mkNonCommTermEqProof (e e' : RingExpr) : NonCommRingM Expr := do
+  let ring ← getRing
+  let { lhs, lhs', vars, .. } := norm ring.vars e (.num 0) e' (.num 0)
+  let ctx ← toContextExpr vars
+  let h := mkApp2 (mkConst ``Grind.CommRing.Expr.eq_of_toPoly_nc_eq [ring.u]) ring.type ring.ringInst
+  return mkApp4 h ctx (toExpr lhs) (toExpr lhs') eagerReflBoolTrue
+
 end Lean.Meta.Grind.Arith.CommRing

src/Lean/Meta/Tactic/Grind/Order/Assert.lean

@@ -75,7 +75,7 @@ public def propagateEqTrue (c : Cnstr NodeId) (e : Expr) (u v : NodeId) (k k' :
   let mut h ← mkPropagateEqTrueProof u v k kuv k'
   if let some he := c.h? then
     h := mkApp4 (mkConst ``Grind.Order.eq_trans_true) e c.e he h
-  if let some (e', he) := (← get').cnstrsMapInv.find? { expr := e } then
+  if let some (e', he) := (← get').termMapInv.find? { expr := e } then
     h := mkApp4 (mkConst ``Grind.Order.eq_trans_true) e' e he h
     pushEqTrue e' h
   else
@@ -87,7 +87,7 @@ public def propagateSelfEqTrue (c : Cnstr NodeId) (e : Expr) : OrderM Unit := do
   let mut h ← mkPropagateSelfEqTrueProof u c.getWeight
   if let some he := c.h? then
     h := mkApp4 (mkConst ``Grind.Order.eq_trans_true) e c.e he h
-  if let some (e', he) := (← get').cnstrsMapInv.find? { expr := e } then
+  if let some (e', he) := (← get').termMapInv.find? { expr := e } then
     h := mkApp4 (mkConst ``Grind.Order.eq_trans_true) e' e he h
     pushEqTrue e' h
   else
@@ -100,7 +100,7 @@ public def propagateEqFalse (c : Cnstr NodeId) (e : Expr) (u v : NodeId) (k k' :
   let mut h ← mkPropagateEqFalseProof u v k kuv k'
   if let some he := c.h? then
     h := mkApp4 (mkConst ``Grind.Order.eq_trans_false) e c.e he h
-  if let some (e', he) := (← get').cnstrsMapInv.find? { expr := e } then
+  if let some (e', he) := (← get').termMapInv.find? { expr := e } then
     h := mkApp4 (mkConst ``Grind.Order.eq_trans_false) e' e he h
     pushEqFalse e' h
   else
@@ -112,7 +112,7 @@ public def propagateSelfEqFalse (c : Cnstr NodeId) (e : Expr) : OrderM Unit := d
   let mut h ← mkPropagateSelfEqFalseProof u c.getWeight
   if let some he := c.h? then
     h := mkApp4 (mkConst ``Grind.Order.eq_trans_false) e c.e he h
-  if let some (e', he) := (← get').cnstrsMapInv.find? { expr := e } then
+  if let some (e', he) := (← get').termMapInv.find? { expr := e } then
     h := mkApp4 (mkConst ``Grind.Order.eq_trans_false) e' e he h
     pushEqFalse e' h
   else
@@ -140,7 +140,7 @@ Returns `true` if `e` is already `True` in the `grind` core.
 Recall that `e` may be an auxiliary term created for a term `e'` (see `cnstrsMapInv`).
 -/
 private def isAlreadyTrue (e : Expr) : OrderM Bool := do
-  if let some (e', _) := (← get').cnstrsMapInv.find? { expr := e } then
+  if let some (e', _) := (← get').termMapInv.find? { expr := e } then
     alreadyInternalized e' <&&> isEqTrue e'
   else
     alreadyInternalized e <&&> isEqTrue e
@@ -162,7 +162,7 @@ Returns `true` if `e` is already `False` in the `grind` core.
 Recall that `e` may be an auxiliary term created for a term `e'` (see `cnstrsMapInv`).
 -/
 private def isAlreadyFalse (e : Expr) : OrderM Bool := do
-  if let some (e', _) := (← get').cnstrsMapInv.find? { expr := e } then
+  if let some (e', _) := (← get').termMapInv.find? { expr := e } then
     alreadyInternalized e' <&&> isEqFalse e'
   else
     alreadyInternalized e <&&> isEqFalse e
@@ -221,7 +221,7 @@ Adds an edge `u --(k) --> v` justified by the proof term `p`, and then
 if no negative cycle was created, updates the shortest distance of affected
 node pairs.
 -/
-def addEdge (u : NodeId) (v : NodeId) (k : Weight) (h : Expr) : OrderM Unit := do
+public def addEdge (u : NodeId) (v : NodeId) (k : Weight) (h : Expr) : OrderM Unit := do
   if (← isInconsistent) then return ()
   if u == v then
     if k.isNeg then
@@ -303,7 +303,7 @@ def getStructIdOf? (e : Expr) : GoalM (Option Nat) := do
   return (← get').exprToStructId.find? { expr := e }
 
 def propagateIneq (e : Expr) : GoalM Unit := do
-  if let some (e', he) := (← get').cnstrsMap.find? { expr := e } then
+  if let some (e', he) := (← get').termMap.find? { expr := e } then
     go e' (some he)
   else
     go e none

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

@@ -7,6 +7,8 @@ module
 prelude
 public import Lean.Meta.Tactic.Grind.Order.OrderM
 import Init.Data.Int.OfNat
+import Init.Grind.Module.Envelope
+import Init.Grind.Order
 import Lean.Meta.Tactic.Grind.Arith.CommRing.SafePoly
 import Lean.Meta.Tactic.Grind.Arith.CommRing.Reify
 import Lean.Meta.Tactic.Grind.Arith.CommRing.DenoteExpr
@@ -15,6 +17,7 @@ import Lean.Meta.Tactic.Grind.Arith.Cutsat.Nat
 import Lean.Meta.Tactic.Grind.Order.StructId
 import Lean.Meta.Tactic.Grind.Order.Util
 import Lean.Meta.Tactic.Grind.Order.Assert
+import Lean.Meta.Tactic.Grind.Order.Proof
 namespace Lean.Meta.Grind.Order
 
 open Arith CommRing
@@ -40,8 +43,24 @@ def getType? (e : Expr) : Option Expr :=
   | _ => none
 
 def isForbiddenParent (parent? : Option Expr) : Bool :=
-  if let some parent := parent? then
-    getType? parent |>.isSome
+  if isIntModuleVirtualParent parent? then
+    /-
+    **Note**: `linarith` uses a virtual parent to mark auxiliary declarations used to encode
+    terms into an `IntModule`.
+    -/
+    true
+  else if let some parent := parent? then
+    (getType? parent |>.isSome)
+    ||
+    /-
+    **Note**: We currently ignore `•`. We may reconsider it in the future.
+    -/
+    match_expr parent with
+    | HSMul.hSMul _ _ _ _ _ _ => true
+    | Nat.cast _ _ _ => true
+    | NatCast.natCast _ _ _ => true
+    | Grind.IntModule.OfNatModule.toQ _ _ _ => true
+    | _ => false
   else
     false
 
@@ -206,27 +225,89 @@ def internalizeCnstr (e : Expr) (kind : CnstrKind) (lhs rhs : Expr) : OrderM Uni
       s.cnstrsOf.insert (u, v) cs
   }
 
+/--
+Normalization result. A nested term `e` is normalized as `a + k` and
+`h` is a proof for `e = a + k`
+-/
+structure OffsetTermResult where
+  a : Expr
+  k : Int
+  h : Expr
+
+def toOffsetTermCommRing? (e : Expr) : RingM (Option OffsetTermResult) := do
+  let some e ← reify? e (skipVar := false) | return none
+  let p ← e.toPolyM
+  let k := p.getConst
+  let p := p.addConst (-k)
+  let a ← shareCommon (← p.denoteExpr)
+  let h ← mkTermEqProof e (.add (p.toExpr) (.intCast k))
+  return some { a, k, h }
+
+def toOffsetTermNonCommRing? (e : Expr) : NonCommRingM (Option OffsetTermResult) := do
+  let some e ← ncreify? e (skipVar := false) | return none
+  let p := e.toPoly_nc
+  let k := p.getConst
+  let p := p.addConst (-k)
+  let a ← shareCommon (← p.denoteExpr)
+  let h ← mkNonCommTermEqProof e (.add (p.toExpr) (.intCast k))
+  return some { a, k, h }
+
+def toOffsetTerm? (e : Expr) : OrderM (Option OffsetTermResult) := do
+  let s ← getStruct
+  /-
+  **Note**: If it is not a partial order, then it is not worth internalizing term
+  since we will not be able to propagate implied equalities back to core.
+  -/
+  unless s.isPartialInst?.isSome do return none
+  let some ringId := s.ringId? | return none
+  if s.isCommRing then
+    RingM.run ringId <| toOffsetTermCommRing? e
+  else
+    NonCommRingM.run ringId <| toOffsetTermNonCommRing? e
+
+def internalizeTerm (e : Expr) : OrderM Unit := do
+  let some r ← toOffsetTerm? e | return ()
+  let x ← mkNode e
+  let y ← mkNode r.a
+  let h₁ ← mkOrdRingPrefix ``Grind.Order.le_of_offset_eq_1_k
+  let h₁ := mkApp4 h₁ e r.a (toExpr r.k) r.h
+  addEdge x y { k := r.k } h₁
+  let h₂ ← mkOrdRingPrefix ``Grind.Order.le_of_offset_eq_2_k
+  let h₂ := mkApp4 h₂ e r.a (toExpr r.k) r.h
+  addEdge y x { k := -r.k } h₂
+
+def updateTermMap (e eNew h : Expr) : GoalM Unit := do
+  modify' fun s => { s with
+    termMap    := s.termMap.insert { expr := e } (eNew, h)
+    termMapInv := s.termMapInv.insert { expr := eNew } (e, h)
+  }
+
 open Arith.Cutsat in
 def adaptNat (e : Expr) : GoalM Expr := do
-  let (eNew, h) ← match_expr e with
-    | LE.le _ _ lhs rhs =>
-      let (lhs', h₁) ← natToInt lhs
-      let (rhs', h₂) ← natToInt rhs
-      let eNew := mkIntLE lhs' rhs'
-      let h := mkApp6 (mkConst ``Nat.ToInt.le_eq) lhs rhs lhs' rhs' h₁ h₂
-      pure (eNew, h)
-    | LT.lt _ _ lhs rhs =>
-      let (lhs', h₁) ← natToInt lhs
-      let (rhs', h₂) ← natToInt rhs
-      let eNew := mkIntLT lhs' rhs'
-      let h := mkApp6 (mkConst ``Nat.ToInt.lt_eq) lhs rhs lhs' rhs' h₁ h₂
-      pure (eNew, h)
+  if let some (eNew, _) := (← get').termMap.find? { expr := e } then
+    return eNew
+  else match_expr e with
+    | LE.le _ _ lhs rhs => adaptCnstr lhs rhs (isLT := false)
+    | LT.lt _ _ lhs rhs => adaptCnstr lhs rhs (isLT := true)
+    | HAdd.hAdd _ _ _ _ _ _ => adaptTerm
+    | HSub.hSub _ _ _ _ _ _ => adaptTerm
+    | OfNat.ofNat _ _ _ => adaptTerm
     | _ => return e
-  modify' fun s => { s with
-    cnstrsMap    := s.cnstrsMap.insert { expr := e } (eNew, h)
-    cnstrsMapInv := s.cnstrsMapInv.insert { expr := eNew } (e, h)
-  }
-  return eNew
+where
+  adaptCnstr (lhs rhs : Expr) (isLT : Bool) : GoalM Expr := do
+    let (lhs', h₁) ← natToInt lhs
+    let (rhs', h₂) ← natToInt rhs
+    let eNew := if isLT then mkIntLT lhs' rhs' else mkIntLE lhs' rhs'
+    let h := mkApp6
+        (mkConst (if isLT then ``Nat.ToInt.lt_eq else ``Nat.ToInt.le_eq))
+        lhs rhs lhs' rhs' h₁ h₂
+    updateTermMap e eNew h
+    return eNew
+
+  adaptTerm : GoalM Expr := do
+    let (eNew, h) ← natToInt e
+    updateTermMap e eNew h
+    return eNew
 
 def adapt (α : Expr) (e : Expr) : GoalM (Expr × Expr) := do
   -- **Note**: We currently only adapt `Nat` expressions
@@ -250,8 +331,9 @@ public def internalize (e : Expr) (parent? : Option Expr) : GoalM Unit := do
     match_expr e with
     | LE.le _ _ lhs rhs => internalizeCnstr e .le lhs rhs
     | LT.lt _ _ lhs rhs => if (← hasLt) then internalizeCnstr e .lt lhs rhs
-    | _ =>
-      -- **Note**: We currently do not internalize offset terms nested in other terms.
-      return ()
+    | HAdd.hAdd _ _ _ _ _ _ => internalizeTerm e
+    | HSub.hSub _ _ _ _ _ _ => internalizeTerm e
+    | OfNat.ofNat _ _ _ => internalizeTerm e
+    | _ => return ()
 
 end Lean.Meta.Grind.Order

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

@@ -144,9 +144,9 @@ structure State where
   Example: given `x y : Nat`, `x ≤ y + 1` is mapped to `Int.ofNat x ≤ Int.ofNat y + 1`, and proof
   of equivalence.
   -/
-  cnstrsMap    : PHashMap ExprPtr (Expr × Expr) := {}
-  /-- `cnstrsMap` inverse -/
-  cnstrsMapInv : PHashMap ExprPtr (Expr × Expr) := {}
+  termMap    : PHashMap ExprPtr (Expr × Expr) := {}
+  /-- `termMap` inverse -/
+  termMapInv : PHashMap ExprPtr (Expr × Expr) := {}
   deriving Inhabited
 
 builtin_initialize orderExt : SolverExtension State ← registerSolverExtension (return {})

tests/lean/run/grind_cutsat_toint_1.lean

@@ -90,7 +90,7 @@ example (a b : Fin 3) : a > 0 → a ≠ b → a + b ≠ 0 → a + b ≠ 1 → Fa
   grind
 
 -- We use `↑a` when pretty printing `ToInt.toInt a`
-/-- trace: [grind.debug.ring.basis] (↑a + ↑b) % 3 + -1 * ↑a + -1 * ↑b + 3 * ((↑a + ↑b) / 3) = 0 -/
+/-- trace: [grind.debug.ring.basis] ↑a + ↑b + -1 * ((↑a + ↑b) % 3) + -3 * ((↑a + ↑b) / 3) = 0 -/
 #guard_msgs (drop error, trace) in
 set_option trace.grind.debug.ring.basis true in
 example (a b : Fin 3) : a > 0 → a ≠ b → a + b ≠ 0 → a + b ≠ 1 → False := by

tests/lean/run/grind_order_eq.lean

@@ -0,0 +1,11 @@
+open Lean Grind Std
+
+example [LE α] [IsPartialOrder α] (a b : α) (f : α → Nat) : a ≤ b → b ≤ c → c ≤ a → f a = f b := by
+  grind (splits := 0)
+
+example [CommRing α] [LE α] [LT α] [LawfulOrderLT α] [IsPartialOrder α] [OrderedRing α]
+    (a b : α) (f : α → Int) : a ≤ b + 1 → b ≤ a - 1 → f a = f (2 + b - 1) := by
+  grind -mbtc -lia -linarith (splits := 0)
+
+example (a b : Int) (f : Int → Int) : a ≤ b + 1 → b ≤ a - 1 → f a = f (2 + b - 1) := by
+  grind -mbtc -lia -linarith (splits := 0)