feat: use CommRing to normalize linarith expressions (#8682)

This PR uses the `CommRing` module to normalize linarith inequalities.

Author: leodemoura

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

@@ -31,8 +31,11 @@ private def reportAppIssue (e : Expr) : GoalM Unit := do
 /--
 Converts a Lean expression `e` in the `CommRing` with id `ringId` into
 a `CommRing.Expr` object.
+
+If `skipVar` is `true`, then the result is `none` if `e` is not an interpreted `CommRing` term.
+We use `skipVar := false` when processing inequalities, and `skipVar := true` for equalities and disequalities
 -/
-partial def reify? (e : Expr) : RingM (Option RingExpr) := do
+partial def reify? (e : Expr) (skipVar := true) : RingM (Option RingExpr) := do
   let toVar (e : Expr) : RingM RingExpr := do
     return .var (← mkVar e)
   let asVar (e : Expr) : RingM RingExpr := do
@@ -67,36 +70,41 @@ partial def reify? (e : Expr) : RingM (Option RingExpr) := do
       let some k ← getNatValue? n | toVar e
       return .num k
     | _ => toVar e
-  let asNone (e : Expr) : GoalM (Option RingExpr) := do
+  let toTopVar (e : Expr) : RingM (Option RingExpr) := do
+    if skipVar then
+      return none
+    else
+      return some (← toVar e)
+  let asTopVar (e : Expr) : RingM (Option RingExpr) := do
     reportAppIssue e
-    return none
+    toTopVar e
   match_expr e with
   | HAdd.hAdd _ _ _ i a b =>
-    if isAddInst (← getRing) i then return some (.add (← go a) (← go b)) else asNone e
+    if isAddInst (← getRing) i then return some (.add (← go a) (← go b)) else asTopVar e
   | HMul.hMul _ _ _ i a b =>
-    if isMulInst (← getRing) i then return some (.mul (← go a) (← go b)) else asNone e
+    if isMulInst (← getRing) i then return some (.mul (← go a) (← go b)) else asTopVar e
   | HSub.hSub _ _ _ i a b =>
-    if isSubInst (← getRing) i then return some (.sub (← go a) (← go b)) else asNone e
+    if isSubInst (← getRing) i then return some (.sub (← go a) (← go b)) else asTopVar e
   | HPow.hPow _ _ _ i a b =>
     let some k ← getNatValue? b | return none
-    if isPowInst (← getRing) i then return some (.pow (← go a) k) else asNone e
+    if isPowInst (← getRing) i then return some (.pow (← go a) k) else asTopVar e
   | Neg.neg _ i a =>
-    if isNegInst (← getRing) i then return some (.neg (← go a)) else asNone e
-  | IntCast.intCast _ i e =>
+    if isNegInst (← getRing) i then return some (.neg (← go a)) else asTopVar e
+  | IntCast.intCast _ i a =>
     if isIntCastInst (← getRing) i then
-      let some k ← getIntValue? e | return none
+      let some k ← getIntValue? a | asTopVar e
       return some (.num k)
     else
-      asNone e
-  | NatCast.natCast _ i e =>
+      asTopVar e
+  | NatCast.natCast _ i a =>
     if isNatCastInst (← getRing) i then
-      let some k ← getNatValue? e | return none
+      let some k ← getNatValue? a | asTopVar e
       return some (.num k)
     else
-      asNone e
+      asTopVar e
   | OfNat.ofNat _ n _ =>
-    let some k ← getNatValue? n | return none
+    let some k ← getNatValue? n | asTopVar e
     return some (.num k)
-  | _ => return none
+  | _ => toTopVar e
 
 end  Lean.Meta.Grind.Arith.CommRing

src/Lean/Meta/Tactic/Grind/Arith/Linear/DenoteExpr.lean

@@ -46,21 +46,58 @@ private def mkEq (a b : Expr) : M Expr := do
   return mkApp3 (mkConst ``Eq [s.u.succ]) s.type a b
 
 def EqCnstr.denoteExpr (c : EqCnstr) : M Expr := do
-  mkEq (← c.p.denoteExpr) (← getStruct).zero
+  mkEq (← c.p.denoteExpr) (← getStruct).ofNatZero
 
 def DiseqCnstr.denoteExpr (c : DiseqCnstr) : M Expr := do
-  return mkNot (← mkEq (← c.p.denoteExpr) (← getStruct).zero)
+  return mkNot (← mkEq (← c.p.denoteExpr) (← getStruct).ofNatZero)
 
 private def denoteIneq (p : Poly) (strict : Bool) : M Expr := do
   if strict then
-    return mkApp2 (← getStruct).ltFn (← p.denoteExpr) (← getStruct).zero
+    return mkApp2 (← getStruct).ltFn (← p.denoteExpr) (← getStruct).ofNatZero
   else
-    return mkApp2 (← getStruct).leFn (← p.denoteExpr) (← getStruct).zero
+    return mkApp2 (← getStruct).leFn (← p.denoteExpr) (← getStruct).ofNatZero
 
 def IneqCnstr.denoteExpr (c : IneqCnstr) : M Expr := do
   denoteIneq c.p c.strict
 
 def NotIneqCnstr.denoteExpr (c : NotIneqCnstr) : M Expr := do
   return mkNot (← denoteIneq c.p c.strict)
 
+private def denoteNum (k : Int) : LinearM Expr := do
+  return mkApp2 (← getStruct).hmulFn (mkIntLit k) (← getOne)
+
+def _root_.Lean.Grind.CommRing.Power.denoteAsIntModuleExpr (pw : Grind.CommRing.Power) : LinearM Expr := do
+  let x := (← getRing).vars[pw.x]!
+  if pw.k == 1 then
+    return x
+  else
+    return mkApp2 (← getRing).powFn x (toExpr pw.k)
+
+def _root_.Lean.Grind.CommRing.Mon.denoteAsIntModuleExpr (m : Grind.CommRing.Mon) : LinearM Expr := do
+  match m with
+  | .unit => getOne
+  | .mult pw m => go m (← pw.denoteAsIntModuleExpr)
+where
+  go (m : Grind.CommRing.Mon) (acc : Expr) : LinearM Expr := do
+    match m with
+    | .unit => return acc
+    | .mult pw m => go m (mkApp2 (← getRing).mulFn acc (← pw.denoteAsIntModuleExpr))
+
+def _root_.Lean.Grind.CommRing.Poly.denoteAsIntModuleExpr (p : Grind.CommRing.Poly) : LinearM Expr := do
+  match p with
+  | .num k => denoteNum k
+  | .add k m p => go p (← denoteTerm k m)
+where
+  denoteTerm (k : Int) (m : Grind.CommRing.Mon) : LinearM Expr := do
+    if k == 1 then
+      m.denoteAsIntModuleExpr
+    else
+      return mkApp2 (← getStruct).hmulFn (mkIntLit k) (← m.denoteAsIntModuleExpr)
+
+  go (p : Grind.CommRing.Poly) (acc : Expr) : LinearM Expr := do
+    match p with
+    | .num 0 => return acc
+    | .num k => return mkApp2 (← getStruct).addFn acc (← denoteNum k)
+    | .add k m p => go p (mkApp2 (← getStruct).addFn acc (← denoteTerm k m))
+
 end Lean.Meta.Grind.Arith.Linear

src/Lean/Meta/Tactic/Grind/Arith/Linear/IneqCnstr.lean

@@ -4,6 +4,9 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 prelude
+import Init.Grind.CommRing.Poly
+import Lean.Meta.Tactic.Grind.Arith.CommRing.Reify
+import Lean.Meta.Tactic.Grind.Arith.CommRing.DenoteExpr
 import Lean.Meta.Tactic.Grind.Arith.Linear.Var
 import Lean.Meta.Tactic.Grind.Arith.Linear.StructId
 import Lean.Meta.Tactic.Grind.Arith.Linear.Reify
@@ -16,6 +19,57 @@ def isLeInst (struct : Struct) (inst : Expr) : Bool :=
 def isLtInst (struct : Struct) (inst : Expr) : Bool :=
   isSameExpr struct.ltFn.appArg! inst
 
+def IneqCnstr.assert (c : IneqCnstr) : LinearM Unit := do
+  trace[grind.linarith.assert] "{← c.denoteExpr}"
+  -- TODO
+
+def NotIneqCnstr.assert (c : NotIneqCnstr) : LinearM Unit := do
+  trace[grind.linarith.assert] "{← c.denoteExpr}"
+  -- TODO
+
+def propagateCommRingIneq (e : Expr) (lhs rhs : Expr) (strict : Bool) (eqTrue : Bool) : LinearM Unit := do
+  let some lhs ← withRingM <| CommRing.reify? lhs (skipVar := false) | return ()
+  let some rhs ← withRingM <| CommRing.reify? rhs (skipVar := false) | return ()
+  if eqTrue then
+    let p' := (lhs.sub rhs).toPoly
+    let lhs' ← p'.denoteAsIntModuleExpr
+    let some lhs' ← reify? lhs' (skipVar := false) | return ()
+    let p := lhs'.norm
+    let c : IneqCnstr := { p, strict, h := .coreCommRing e lhs rhs lhs' }
+    c.assert
+  else if (← isLinearOrder) then
+    let p' := (rhs.sub lhs).toPoly
+    let strict := !strict
+    let lhs' ← p'.denoteAsIntModuleExpr
+    let some lhs' ← reify? lhs' (skipVar := false) | return ()
+    let p := lhs'.norm
+    let c : IneqCnstr := { p, strict, h := .notCoreCommRing e lhs rhs lhs' }
+    c.assert
+  else
+    let p' := (lhs.sub rhs).toPoly
+    let lhs' ← p'.denoteAsIntModuleExpr
+    let some lhs' ← reify? lhs' (skipVar := false) | return ()
+    let p := lhs'.norm
+    let c : NotIneqCnstr := { p, strict, h := .coreCommRing e lhs rhs lhs' }
+    c.assert
+
+def propagateIntModuleIneq (e : Expr) (lhs rhs : Expr) (strict : Bool) (eqTrue : Bool) : LinearM Unit := do
+  let some lhs ← reify? lhs (skipVar := false) | return ()
+  let some rhs ← reify? rhs (skipVar := false) | return ()
+  if eqTrue then
+    let p := (lhs.sub rhs).norm
+    let c : IneqCnstr := { p, strict, h := .core e lhs rhs }
+    c.assert
+  else if (← isLinearOrder) then
+    let p := (rhs.sub lhs).norm
+    let strict := !strict
+    let c : IneqCnstr := { p, strict, h := .notCore e lhs rhs }
+    c.assert
+  else
+    let p := (lhs.sub rhs).norm
+    let c : NotIneqCnstr := { p, strict, h := .core e lhs rhs }
+    c.assert
+
 def propagateIneq (e : Expr) (eqTrue : Bool) : GoalM Unit := do
   let numArgs := e.getAppNumArgs
   unless numArgs == 4 do return ()
@@ -29,15 +83,13 @@ def propagateIneq (e : Expr) (eqTrue : Bool) : GoalM Unit := do
     else if isLtInst struct inst then
       pure true
     else
-      trace[grind.linarith] "invalid {e}, {(← getStruct).leFn}, {(← getStruct).ltFn}"
       return ()
-    let some lhs ← reify? (e.getArg! 2 numArgs) (skipVar := false) | trace[grind.linarith] "lhs failed {e}"; return ()
-    let some rhs ← reify? (e.getArg! 3 numArgs) (skipVar := false) | trace[grind.linarith] "rhs failed {e}"; return ()
-    let p := (lhs.sub rhs).norm
-    -- TODO
-    trace[grind.linarith] "{e}, {eqTrue}, strict: {strict}, structId: {structId}"
-    trace[grind.linarith] "{← p.denoteExpr}"
-    trace[grind.linarith] "structId: {structId}"
-    return ()
+    let lhs := e.getArg! 2 numArgs
+    let rhs := e.getArg! 3 numArgs
+    if (← isCommRing) then
+      propagateCommRingIneq e lhs rhs strict eqTrue
+    -- TODO: non-commutative ring normalizer
+    else
+      propagateIntModuleIneq e lhs rhs strict eqTrue
 
 end Lean.Meta.Grind.Arith.Linear

src/Lean/Meta/Tactic/Grind/Arith/Linear/Reify.lean

@@ -41,10 +41,7 @@ partial def reify? (e : Expr) (skipVar : Bool) : LinearM (Option LinExpr) := do
     reportInstIssue e
     return .var (← mkVar e)
   let isOfNatZero (e : Expr) : LinearM Bool := do
-    let_expr OfNat.ofNat _ n _ := e | return false
-    let some k ← getNatValue? n | return false
-    unless k == 0 do return false
-    withDefault <| isDefEq e (← getStruct).zero
+    withDefault <| isDefEq e (← getStruct).ofNatZero
   let rec go (e : Expr) : LinearM   LinExpr := do
     match_expr e with
     | HAdd.hAdd _ _ _ i a b =>

src/Lean/Meta/Tactic/Grind/Arith/Linear/StructId.lean

@@ -7,13 +7,22 @@ prelude
 import Init.Grind.Ordered.Module
 import Lean.Meta.Tactic.Grind.Simp
 import Lean.Meta.Tactic.Grind.Internalize
+import Lean.Meta.Tactic.Grind.Arith.CommRing.RingId
 import Lean.Meta.Tactic.Grind.Arith.Linear.Util
 import Lean.Meta.Tactic.Grind.Arith.Linear.Var
 
 namespace Lean.Meta.Grind.Arith.Linear
 
+private def preprocess (e : Expr) : GoalM Expr := do
+  shareCommon (← canon e)
+
 private def internalizeFn (fn : Expr) : GoalM Expr := do
-  shareCommon (← canon fn)
+  preprocess fn
+
+private def preprocessConst (c : Expr) : GoalM Expr := do
+  let c ← preprocess c
+  internalize c none
+  return c
 
 private def internalizeConst (c : Expr) : GoalM Expr := do
   let c ← shareCommon (← canon c)
@@ -22,6 +31,13 @@ private def internalizeConst (c : Expr) : GoalM Expr := do
 
 open Grind.Linarith (Poly)
 
+private def mkExpectedDefEqMsg (a b : Expr) : MetaM MessageData :=
+  return m!"`grind linarith` expected{indentExpr a}\nto be definitionally equal to{indentExpr b}"
+
+private def ensureDefEq (a b : Expr) : MetaM Unit := do
+  unless (← withDefault <| isDefEq a b) do
+    throwError (← mkExpectedDefEqMsg a b)
+
 def getStructId? (type : Expr) : GoalM (Option Nat) := do
   if let some id? := (← get').typeIdOf.find? { expr := type } then
     return id?
@@ -55,21 +71,24 @@ where
       let some inst := inst? | return none
       let toField := mkApp2 (mkConst toFieldName [u]) type inst
       unless (← withDefault <| isDefEq parentInst toField) do
-        reportIssue! "`grind linarith` expected{indentExpr parentInst}\nto be definitionally equal to{indentExpr toField}"
+        reportIssue! (← mkExpectedDefEqMsg parentInst toField)
         return none
       return some inst
     let ensureToFieldDefEq (parentInst : Expr) (inst : Expr) (toFieldName : Name) : GoalM Unit := do
       let toField := mkApp2 (mkConst toFieldName [u]) type inst
-      unless (← withDefault <| isDefEq parentInst toField) do
-        throwError "`grind linarith` expected{indentExpr parentInst}\nto be definitionally equal to{indentExpr toField}"
+      ensureDefEq parentInst toField
     let ensureToHomoFieldDefEq (parentInst : Expr) (inst : Expr) (toFieldName : Name) (toHeteroName : Name) : GoalM Unit := do
       let toField := mkApp2 (mkConst toFieldName [u]) type inst
       let heteroToField := mkApp2 (mkConst toHeteroName [u]) type toField
-      unless (← withDefault <| isDefEq parentInst heteroToField) do
-        throwError "`grind linarith` expected{indentExpr parentInst}\nto be definitionally equal to{indentExpr heteroToField}"
+      ensureDefEq parentInst heteroToField
     let some intModuleInst ← getInst? ``Grind.IntModule | return none
     let zeroInst ← getInst ``Zero
     let zero ← internalizeConst <| mkApp2 (mkConst ``Zero.zero [u]) type zeroInst
+    let ofNatZeroType := mkApp2 (mkConst ``OfNat [u]) type (mkRawNatLit 0)
+    let some ofNatZeroInst := LOption.toOption (← trySynthInstance ofNatZeroType) | return none
+    -- `ofNatZero` is used internally, we don't need to internalize
+    let ofNatZero ← preprocess <| mkApp3 (mkConst ``OfNat.ofNat [u]) type (mkRawNatLit 0) ofNatZeroInst
+    ensureDefEq zero ofNatZero
     let addInst ← getBinHomoInst ``HAdd
     let addFn ← internalizeFn <| mkApp4 (mkConst ``HAdd.hAdd [u, u, u]) type type type addInst
     let subInst ← getBinHomoInst ``HSub
@@ -100,15 +119,16 @@ where
       let smulFn ← internalizeFn <| mkApp4 (mkConst ``HSMul.hSMul [0, u, u]) Int.mkType type smulInst smulInst
       if (← withDefault <| isDefEq hmulFn smulFn) then
         return smulFn
-      reportIssue! "`grind linarith` expected{indentExpr hmulFn}\nto be definitionally equal to{indentExpr smulFn}"
+      reportIssue! (← mkExpectedDefEqMsg hmulFn smulFn)
       return none
     let smulFn? ← getSMulFn?
+    let ringId? ← CommRing.getRingId? type
     let ringInst? ← getInst? ``Grind.Ring
     let getOne? : GoalM (Option Expr) := do
       let some oneInst ← getInst? ``One | return none
       let one ← internalizeConst <| mkApp2 (mkConst ``One.one [u]) type oneInst
       let one' ← mkNumeral type 1
-      unless (← withDefault <| isDefEq one one') do reportIssue! "`grind linarith` expected{indentExpr one}\nto be definitionally equal to{indentExpr one'}"
+      unless (← withDefault <| isDefEq one one') do reportIssue! (← mkExpectedDefEqMsg one one')
       return some one
     let one? ← getOne?
     let commRingInst? ← getInst? ``Grind.CommRing
@@ -127,7 +147,7 @@ where
     let struct : Struct := {
       id, type, u, intModuleInst, preorderInst, isOrdInst, partialInst?, linearInst?, noNatDivInst?
       leFn, ltFn, addFn, subFn, negFn, hmulFn, smulFn?, zero, one?
-      ringInst?, commRingInst?, ringIsOrdInst?
+      ringInst?, commRingInst?, ringIsOrdInst?, ringId?, ofNatZero
     }
     modify' fun s => { s with structs := s.structs.push struct }
     if let some one := one? then

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

@@ -5,6 +5,7 @@ Authors: Leonardo de Moura
 -/
 prelude
 import Std.Internal.Rat
+import Init.Grind.CommRing.Poly
 import Init.Grind.Ordered.Linarith
 import Lean.Data.PersistentArray
 import Lean.Meta.Tactic.Grind.ExprPtr
@@ -38,6 +39,8 @@ structure IneqCnstr where
 inductive IneqCnstrProof where
   | core (e : Expr) (lhs rhs : LinExpr)
   | notCore (e : Expr) (lhs rhs : LinExpr)
+  | coreCommRing (e : Expr) (lhs rhs : Grind.CommRing.Expr) (lhs' : LinExpr)
+  | notCoreCommRing (e : Expr) (lhs rhs : Grind.CommRing.Expr) (lhs' : LinExpr)
   | combine (c₁ : IneqCnstr) (c₂ : IneqCnstr)
   | combineEq (c₁ : IneqCnstr) (c₂ : EqCnstr)
   | norm (c₁ : IneqCnstr) (k : Nat)
@@ -62,6 +65,7 @@ structure NotIneqCnstr where
 
 inductive NotIneqCnstrProof where
   | core (e : Expr) (lhs rhs : LinExpr)
+  | coreCommRing (e : Expr) (lhs rhs : Grind.CommRing.Expr) (lhs' : LinExpr)
   -- TODO: norm, and combineEq
 
 inductive UnsatProof where
@@ -77,6 +81,8 @@ Each type must be at least implement the instances `IntModule`, `Preorder`, and
 -/
 structure Struct where
   id               : Nat
+  /-- If the structure is a ring, we store its id in the `CommRing` module at `ringId?` -/
+  ringId?          : Option Nat
   type             : Expr
   /-- Cached `getDecLevel type` -/
   u                : Level
@@ -99,6 +105,7 @@ structure Struct where
   /-- `Ring.IsOrdered` instance with `Preorder` -/
   ringIsOrdInst?   : Option Expr
   zero             : Expr
+  ofNatZero        : Expr
   one?             : Option Expr
   leFn             : Expr
   ltFn             : Expr