test

src/Init/Grind/Module/Envelope.lean

@@ -227,7 +227,7 @@ theorem toQ_add (a b : α) : toQ (a + b) = toQ a + toQ b := by
 theorem toQ_zero : toQ (0 : α) = 0 := by
   simp; apply Quot.sound; simp
 
-theorem toQ_smul (n : Nat) (a : α) : toQ (n • a) = (↑n : Int) • toQ a := by
+theorem toQ_smul (n : Nat) (a : α) : toQ (n • a) = n • toQ a := by
   simp; apply Quot.sound; simp
 
 /-!

src/Init/Grind/Module/OfNatModule.lean

@@ -4,13 +4,11 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-import Init.Grind.Module.Envelope
-
-open Std
-
+public import Init.Grind.Module.Envelope
+public section
 namespace Lean.Grind.IntModule.OfNatModule
+open Std
 
 /-!
 Support for `NatModule` in the `grind` linear arithmetic module.
@@ -44,8 +42,7 @@ theorem add_congr {α} [NatModule α] {a b : α} {a' b' : Q α}
     (h₁ : toQ a = a') (h₂ : toQ b = b') : toQ (a + b) = a' + b' := by
   rw [toQ_add, h₁, h₂]
 
-theorem smul_congr {α} [NatModule α] (n : Nat) (a : α) (i : Int) (a' : Q α)
-    (h₁ : ↑n == i) (h₂ : toQ a = a') : toQ (n • a) = i • a' := by
-  simp at h₁; rw [← h₁, ← h₂, toQ_smul]
+theorem smul_congr {α} [NatModule α] (n : Nat) (a : α) (a' : Q α) (h : toQ a = a') : toQ (n • a) = n • a' := by
+  rw [← h, toQ_smul]
 
 end Lean.Grind.IntModule.OfNatModule

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

@@ -4,7 +4,6 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Meta.Tactic.Grind.Arith.Linear.Types
 public import Lean.Meta.Tactic.Grind.Arith.Linear.Util
@@ -23,9 +22,8 @@ public import Lean.Meta.Tactic.Grind.Arith.Linear.Model
 public import Lean.Meta.Tactic.Grind.Arith.Linear.PP
 public import Lean.Meta.Tactic.Grind.Arith.Linear.MBTC
 public import Lean.Meta.Tactic.Grind.Arith.Linear.VarRename
-
+public import Lean.Meta.Tactic.Grind.Arith.Linear.OfNatModule
 public section
-
 namespace Lean
 
 builtin_initialize registerTraceClass `grind.linarith

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

@@ -6,7 +6,8 @@ Authors: Leonardo de Moura
 module
 prelude
 public import Lean.Meta.Tactic.Grind.Types
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Util
+public import Lean.Meta.Tactic.Grind.Arith.Linear.LinearM
+import Lean.Meta.Tactic.Grind.Arith.Linear.Util
 import Lean.Meta.Tactic.Grind.Simp
 import Lean.Meta.Tactic.Grind.Arith.CommRing.DenoteExpr
 import Lean.Meta.Tactic.Grind.Arith.Linear.Var
@@ -16,7 +17,6 @@ namespace Lean.Meta.Grind.Arith.Linear
 /-!
 Helper functions for converting reified terms back into their denotations.
 -/
-
 variable [Monad M] [MonadGetStruct M] [MonadError M]
 
 def _root_.Lean.Grind.Linarith.Poly.denoteExpr (p : Poly) : M Expr := do

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

@@ -4,19 +4,18 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-public import Init.Grind.Ring.Poly
-public import Lean.Meta.Tactic.Grind.Arith.CommRing.Reify
-public import Lean.Meta.Tactic.Grind.Arith.CommRing.DenoteExpr
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Var
-public import Lean.Meta.Tactic.Grind.Arith.Linear.StructId
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Reify
-public import Lean.Meta.Tactic.Grind.Arith.Linear.DenoteExpr
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Proof
-
+public import Lean.Meta.Tactic.Grind.Arith.Linear.LinearM
+import Init.Grind.Ring.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.Util
+import Lean.Meta.Tactic.Grind.Arith.Linear.StructId
+import Lean.Meta.Tactic.Grind.Arith.Linear.Reify
+import Lean.Meta.Tactic.Grind.Arith.Linear.DenoteExpr
+import Lean.Meta.Tactic.Grind.Arith.Linear.Proof
 public section
-
 namespace Lean.Meta.Grind.Arith.Linear
 
 def isLeInst (struct : Struct) (inst : Expr) : Bool :=

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

@@ -4,19 +4,17 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-public import Lean.Meta.Tactic.Grind.Simp
-public import Lean.Meta.Tactic.Grind.Arith.CommRing.Reify
-public import Lean.Meta.Tactic.Grind.Arith.Linear.StructId
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Reify
-
+public import Lean.Meta.Tactic.Grind.Arith.Linear.LinearM
+public import Lean.Meta.Tactic.Grind.Arith.Linear.OfNatModule
+import Lean.Meta.Tactic.Grind.Simp
+import Lean.Meta.Tactic.Grind.Arith.CommRing.Reify
+import Lean.Meta.Tactic.Grind.Arith.Linear.StructId
+import Lean.Meta.Tactic.Grind.Arith.Linear.Var
+import Lean.Meta.Tactic.Grind.Arith.Linear.Util
+import Lean.Meta.Tactic.Grind.Arith.Linear.Reify
 public section
-
-namespace Lean.Meta.Grind.Arith
-
-
-namespace Linear
+namespace Lean.Meta.Grind.Arith.Linear
 
 /-- If `e` is a function application supported by the linarith module, return its type. -/
 private def getType? (e : Expr) : Option Expr :=
@@ -92,11 +90,13 @@ def internalize (e : Expr) (parent? : Option Expr) : GoalM Unit := do
     return ()
   let some type := getType? e | return ()
   if isForbiddenParent parent? then return ()
-  let some structId ← getStructId? type | return ()
-  LinearM.run structId do
-    trace[grind.linarith.internalize] "{e}"
+  if let some structId ← getStructId? type then LinearM.run structId do
     setTermStructId e
     markAsLinarithTerm e
     markVars e
+  else if let some natStructId ← getNatStructId? type then OfNatModuleM.run natStructId do
+    let (e', _) ← ofNatModule e
+    trace[grind.linarith.internalize] "{e} ==> {e'}"
+    markAsLinarithTerm e
 
 end Lean.Meta.Grind.Arith.Linear

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

@@ -4,14 +4,10 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Util
-
-public section
-
+public import Lean.Meta.Tactic.Grind.Arith.Linear.LinearM
+import Lean.Meta.Tactic.Grind.Arith.Linear.Util
 namespace Lean.Meta.Grind.Arith.Linear
-
 /--
 Returns `true` if the variables in the given polynomial are sorted
 in decreasing order.
@@ -100,7 +96,7 @@ def checkStructInvs : LinearM Unit := do
   checkUppers
   checkDiseqCnstrs
 
-def checkInvariants : GoalM Unit := do
+public def checkInvariants : GoalM Unit := do
   unless grind.debug.get (← getOptions) do return ()
   for structId in *...(← get').structs.size do
     LinearM.run structId do

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

@@ -0,0 +1,88 @@
+/-
+Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+module
+prelude
+public import Lean.Meta.Tactic.Grind.Types
+public import Lean.Meta.Tactic.Grind.Arith.CommRing.RingM
+public section
+namespace Lean.Meta.Grind.Arith.Linear
+
+def get' : GoalM State := do
+  return (← get).arith.linear
+
+@[inline] def modify' (f : State → State) : GoalM Unit := do
+  modify fun s => { s with arith.linear := f s.arith.linear }
+
+structure LinearM.Context where
+  structId : Nat
+
+class MonadGetStruct (m : Type → Type) where
+  getStruct : m Struct
+
+export MonadGetStruct (getStruct)
+
+@[always_inline]
+instance (m n) [MonadLift m n] [MonadGetStruct m] : MonadGetStruct n where
+  getStruct    := liftM (getStruct : m Struct)
+
+/-- We don't want to keep carrying the `StructId` around. -/
+abbrev LinearM := ReaderT LinearM.Context GoalM
+
+abbrev LinearM.run (structId : Nat) (x : LinearM α) : GoalM α :=
+  x { structId }
+
+abbrev getStructId : LinearM Nat :=
+  return (← read).structId
+
+protected def LinearM.getStruct : LinearM Struct := do
+  let s ← get'
+  let structId ← getStructId
+  if h : structId < s.structs.size then
+    return s.structs[structId]
+  else
+    throwError "`grind` internal error, invalid structure id"
+
+instance : MonadGetStruct LinearM where
+  getStruct := LinearM.getStruct
+
+open CommRing
+
+def getRingCore? (ringId? : Option Nat) : GoalM (Option Ring) := do
+  let some ringId := ringId? | return none
+  RingM.run ringId do return some (← getRing)
+
+def throwNotRing : LinearM α :=
+  throwError "`grind linarith` internal error, structure is not a ring"
+
+def throwNotCommRing : LinearM α :=
+  throwError "`grind linarith` internal error, structure is not a commutative ring"
+
+def getRing? : LinearM (Option Ring) := do
+  getRingCore? (← getStruct).ringId?
+
+def getRing : LinearM Ring := do
+  let some ring ← getRing?
+    | throwNotCommRing
+  return ring
+
+instance : MonadRing LinearM where
+  getRing := Linear.getRing
+  modifyRing f := do
+    let some ringId := (← getStruct).ringId? | throwNotCommRing
+    RingM.run ringId do modifyRing f
+  canonExpr e := do shareCommon (← canon e)
+  synthInstance? e := Grind.synthInstance? e
+
+def withRingM (x : RingM α) : LinearM α := do
+  let some ringId := (← getStruct).ringId?
+    | throwNotCommRing
+  RingM.run ringId x
+
+@[inline] def modifyStruct (f : Struct → Struct) : LinearM Unit := do
+  let structId ← getStructId
+  modify' fun s => { s with structs := s.structs.modify structId f }
+
+end Lean.Meta.Grind.Arith.Linear

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

@@ -5,12 +5,12 @@ Authors: Leonardo de Moura
 -/
 module
 prelude
-public import Lean.Meta.Tactic.Grind.MBTC
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Model
-public import Lean.Meta.Tactic.Grind.Arith.Linear.PropagateEq
-
+public import Lean.Meta.Tactic.Grind.Types
+import Lean.Meta.Tactic.Grind.MBTC
+import Lean.Meta.Tactic.Grind.Arith.ModelUtil
+import Lean.Meta.Tactic.Grind.Arith.Linear.Model
+import Lean.Meta.Tactic.Grind.Arith.Linear.PropagateEq
 public section
-
 namespace Lean.Meta.Grind.Arith.Linear
 
 private partial def toRatValue? (a : Expr) : Option Rat :=

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

@@ -4,13 +4,10 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Lean.Meta.Tactic.Grind.Types
-public import Lean.Meta.Tactic.Grind.Arith.ModelUtil
-
+import Lean.Meta.Tactic.Grind.Arith.ModelUtil
 public section
-
 namespace Lean.Meta.Grind.Arith.Linear
 
 def getAssignment? (s : Struct) (e : Expr) : Option Rat := Id.run do

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

@@ -0,0 +1,122 @@
+/-
+Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+module
+prelude
+public import Lean.Meta.Tactic.Grind.Arith.Linear.LinearM
+import Lean.Meta.Tactic.Grind.Simp
+import Init.Grind.Module.OfNatModule
+public section
+namespace Lean.Meta.Grind.Arith.Linear
+
+structure OfNatModuleM.Context where
+  natStructId : Nat
+
+abbrev OfNatModuleM := ReaderT OfNatModuleM.Context GoalM
+
+abbrev OfNatModuleM.run (natStructId : Nat) (x : OfNatModuleM α) : GoalM α :=
+  x { natStructId }
+
+abbrev getNatStructId : OfNatModuleM Nat :=
+  return (← read).natStructId
+
+def getNatStruct : OfNatModuleM NatStruct := do
+  let s ← get'
+  let natStructId ← getNatStructId
+  if h : natStructId < s.natStructs.size then
+    return s.natStructs[natStructId]
+  else
+    throwError "`grind` internal error, invalid natStructId"
+
+protected def OfNatModuleM.getStruct : OfNatModuleM Struct := do
+  let ns ← getNatStruct
+  LinearM.run ns.structId getStruct
+
+instance : MonadGetStruct OfNatModuleM where
+  getStruct := OfNatModuleM.getStruct
+
+@[inline] def modifyNatStruct (f : NatStruct → NatStruct) : OfNatModuleM Unit := do
+  let id ← getNatStructId
+  modify' fun s => { s with natStructs := s.natStructs.modify id f }
+
+def getTermNatStructId? (e : Expr) : GoalM (Option Nat) := do
+  return (← get').exprToNatStructId.find? { expr := e }
+
+/-- Returns `some natStructId` if `a` and `b` are elements of the same `NatModule` structure. -/
+def inSameNatStruct? (a b : Expr) : GoalM (Option Nat) := do
+  let some id ← getTermNatStructId? a | return none
+  let some id' ← getTermNatStructId? b | return none
+  unless id == id' do return none -- This can happen when we have heterogeneous equalities
+  return id
+
+def setTermNatStructId (e : Expr) : OfNatModuleM Unit := do
+  let id ← getNatStructId
+  if let some id' ← getTermNatStructId? e then
+    unless id' == id do
+      reportIssue! "expression in two different nat module structures in linarith module{indentExpr e}"
+    return ()
+  modify' fun s => { s with exprToNatStructId := s.exprToNatStructId.insert { expr := e } id }
+
+private def mkOfNatModuleVar (e : Expr) : OfNatModuleM (Expr × Expr) := do
+  let s ← getNatStruct
+  let toQe := mkApp s.toQFn e
+  let h    := mkApp s.rfl_q toQe
+  setTermNatStructId e
+  markAsLinarithTerm e
+  return (toQe, h)
+
+private def isAddInst (natStruct : NatStruct) (inst : Expr) : Bool :=
+  isSameExpr natStruct.addFn.appArg! inst
+private def isZeroInst (natStruct : NatStruct) (inst : Expr) : Bool :=
+  isSameExpr natStruct.zero.appArg! inst
+private def isSMulInst (natStruct : NatStruct) (inst : Expr) : Bool :=
+  isSameExpr natStruct.smulFn.appArg! inst
+
+private partial def ofNatModule' (e : Expr) : OfNatModuleM (Expr × Expr) := do
+  let s ← getStruct
+  let ns ← getNatStruct
+  match_expr e with
+  | HAdd.hAdd _ _ _ i a b =>
+    if isAddInst ns i then
+      let (a', ha) ← ofNatModule' a
+      let (b', hb) ← ofNatModule' b
+      let e' := mkApp2 s.addFn a' b'
+      let h := mkApp8 (mkConst ``Grind.IntModule.OfNatModule.add_congr [ns.u]) ns.type ns.natModuleInst a b a' b' ha hb
+      pure (e', h)
+    else
+      mkOfNatModuleVar e
+  | HSMul.hSMul _ _ _ i a b =>
+    if isSMulInst ns i then
+      let (b', hb) ← ofNatModule' b
+      let e' := mkApp2 s.nsmulFn a b'
+      let h := mkApp6 (mkConst ``Grind.IntModule.OfNatModule.smul_congr [ns.u]) ns.type ns.natModuleInst a b b' hb
+      pure (e', h)
+    else
+      mkOfNatModuleVar e
+  | Zero.zero _ i =>
+    if isZeroInst ns i then
+      let e' := s.zero
+      let h := mkApp2 (mkConst ``Grind.IntModule.OfNatModule.toQ_zero [ns.u]) ns.type ns.natModuleInst
+      pure (e', h)
+    else
+      mkOfNatModuleVar e
+  | _ => mkOfNatModuleVar e
+
+def ofNatModule (e : Expr) : OfNatModuleM (Expr × Expr) := do
+  if let some r := (← getNatStruct).termMap.find? { expr := e } then
+    return r
+  else
+    let (e', h) ← ofNatModule' e
+    let r ← preprocess e'
+    let (e', h) ← if let some proof := r.proof? then
+      pure (r.expr, (← mkEqTrans h proof))
+    else
+      pure (r.expr, h)
+    setTermNatStructId e
+    internalize e' (← getGeneration e)
+    modifyNatStruct fun s => { s with termMap := s.termMap.insert { expr := e } (e', h) }
+    return (e', h)
+
+end Lean.Meta.Grind.Arith.Linear

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

@@ -4,12 +4,10 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Model
-
+public import Lean.Meta.Tactic.Grind.Types
+import Lean.Meta.Tactic.Grind.Arith.Linear.Model
 public section
-
 namespace Lean.Meta.Grind.Arith.Linear
 
 def ppStruct? (goal : Goal) (s : Struct) : MetaM (Option MessageData) := do

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

@@ -5,7 +5,8 @@ Authors: Leonardo de Moura
 -/
 module
 prelude
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Util
+public import Lean.Meta.Tactic.Grind.Arith.Linear.LinearM
+import Init.Grind.Module.OfNatModule
 import Lean.Data.RArray
 import Lean.Meta.Tactic.Grind.Arith.Util
 import Lean.Meta.Tactic.Grind.Arith.Linear.ToExpr
@@ -15,11 +16,11 @@ import Lean.Meta.Tactic.Grind.Diseq
 import Lean.Meta.Tactic.Grind.ProofUtil
 import Lean.Meta.Tactic.Grind.Arith.CommRing.VarRename
 import Lean.Meta.Tactic.Grind.Arith.CommRing.ToExpr
-import Lean.Meta.Tactic.Grind.Arith.Linear.VarRename
 import Lean.Meta.Tactic.Grind.Arith.CommRing.VarRename
-
+import Lean.Meta.Tactic.Grind.Arith.Linear.VarRename
+import Lean.Meta.Tactic.Grind.Arith.Linear.Util
+import Lean.Meta.Tactic.Grind.Arith.Linear.OfNatModule
 public section
-
 namespace Lean.Meta.Grind.Arith.Linear
 
 open CommRing (RingExpr)
@@ -289,6 +290,15 @@ partial def DiseqCnstr.toExprProof (c' : DiseqCnstr) : ProofM Expr := caching c'
     let h' := mkApp5 h' (← mkRingExprDecl lhs) (← mkRingExprDecl rhs) (← mkRingPolyDecl p') eagerReflBoolTrue (← mkDiseqProof a b)
     let h ← mkIntModThmPrefix ``Grind.Linarith.diseq_norm
     return mkApp5 h (← mkExprDecl lhs') (← mkExprDecl .zero) (← mkPolyDecl c'.p) eagerReflBoolTrue h'
+  | .coreOfNat a b natStructId lhs rhs =>
+    let h ← OfNatModuleM.run natStructId do
+      let ns ← getNatStruct
+      let (a', ha) ← ofNatModule a
+      let (b', hb) ← ofNatModule b
+      return mkApp10 (mkConst ``Grind.IntModule.OfNatModule.of_diseq [ns.u]) ns.type ns.natModuleInst ns.addRightCancelInst?.get!
+        a b a' b' ha hb (← mkDiseqProof a b)
+    return mkApp5 (← mkIntModThmPrefix ``Grind.Linarith.diseq_norm)
+      (← mkExprDecl lhs) (← mkExprDecl rhs) (← mkPolyDecl c'.p) eagerReflBoolTrue h
   | .neg c =>
     let h ← mkIntModThmPrefix ``Grind.Linarith.diseq_neg
     return mkApp4 h (← mkPolyDecl c.p) (← mkPolyDecl c'.p) eagerReflBoolTrue (← c.toExprProof)
@@ -358,7 +368,7 @@ partial def IneqCnstr.collectDecVars (c' : IneqCnstr) : CollectDecVarsM Unit :=
 -- Actually, it cannot even contain decision variables in the current implementation.
 partial def DiseqCnstr.collectDecVars (c' : DiseqCnstr) : CollectDecVarsM Unit := do unless (← alreadyVisited c') do
   match c'.h with
-  | .core .. | .coreCommRing .. | .oneNeZero => return ()
+  | .core .. | .coreCommRing .. | .coreOfNat .. | .oneNeZero => return ()
   | .neg c => c.collectDecVars
   | .subst _ _ c₁ c₂ | .subst1 _ c₁ c₂ => c₁.collectDecVars; c₂.collectDecVars
 

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

@@ -4,20 +4,20 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-public import Init.Grind.Ring.Poly
-public import Lean.Meta.Tactic.Grind.Arith.CommRing.Reify
-public import Lean.Meta.Tactic.Grind.Arith.CommRing.DenoteExpr
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Var
-public import Lean.Meta.Tactic.Grind.Arith.Linear.StructId
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Reify
-public import Lean.Meta.Tactic.Grind.Arith.Linear.IneqCnstr
-public import Lean.Meta.Tactic.Grind.Arith.Linear.DenoteExpr
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Proof
-
-public section
-
+public import Lean.Meta.Tactic.Grind.Arith.Linear.LinearM
+import Init.Grind.Ring.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.Util
+import Lean.Meta.Tactic.Grind.Arith.Linear.StructId
+import Lean.Meta.Tactic.Grind.Arith.Linear.Reify
+import Lean.Meta.Tactic.Grind.Arith.Linear.IneqCnstr
+import Lean.Meta.Tactic.Grind.Arith.Linear.DenoteExpr
+import Lean.Meta.Tactic.Grind.Arith.Linear.Proof
+import Lean.Meta.Tactic.Grind.Arith.Linear.OfNatModule
+section
 namespace Lean.Meta.Grind.Arith.Linear
 
 private def _root_.Lean.Grind.Linarith.Poly.substVar (p : Poly) : LinearM (Option (Var × EqCnstr × Poly)) := do
@@ -258,13 +258,28 @@ private def processNewIntModuleDiseq (a b : Expr) : LinearM Unit := do
   let c : DiseqCnstr := { p, h := .core a b lhs rhs }
   c.assert
 
+private def processNewNatModuleDiseq (a b : Expr) : OfNatModuleM Unit := do
+  let ns ← getNatStruct
+  trace[Meta.debug] "{a}, {b}"
+  unless ns.addRightCancelInst?.isSome do return ()
+  let (a', _) ← ofNatModule a
+  let (b', _) ← ofNatModule b
+  trace[Meta.debug] "{a'}, {b'}"
+  LinearM.run ns.structId do
+    let some lhs ← reify? a' (skipVar := false) | return ()
+    let some rhs ← reify? b' (skipVar := false) | return ()
+    let p := (lhs.sub rhs).norm
+    let c : DiseqCnstr := { p, h := .coreOfNat a b ns.id lhs rhs }
+    c.assert
+
 @[export lean_process_linarith_diseq]
 def processNewDiseqImpl (a b : Expr) : GoalM Unit := do
-  let some structId ← inSameStruct? a b | return ()
-  LinearM.run structId do
+  if let some structId ← inSameStruct? a b then LinearM.run structId do
     if (← isCommRing) then
       processNewCommRingDiseq a b
     else
       processNewIntModuleDiseq a b
+  else if let some natStructId ← inSameNatStruct? a b then OfNatModuleM.run natStructId do
+    processNewNatModuleDiseq a b
 
 end Lean.Meta.Grind.Arith.Linear

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

@@ -4,13 +4,11 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-public import Lean.Meta.Tactic.Grind.Simp
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Var
-
+public import Lean.Meta.Tactic.Grind.Arith.Linear.LinearM
+import Lean.Meta.Tactic.Grind.Simp
+import Lean.Meta.Tactic.Grind.Arith.Linear.Var
 public section
-
 namespace Lean.Meta.Grind.Arith.Linear
 
 def isAddInst (struct : Struct) (inst : Expr) : Bool :=

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

@@ -4,14 +4,13 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-public import Lean.Meta.Tactic.Grind.Arith.Linear.DenoteExpr
 public import Lean.Meta.Tactic.Grind.Arith.Linear.SearchM
-public import Lean.Meta.Tactic.Grind.Arith.Linear.IneqCnstr
-
+import Lean.Meta.Tactic.Grind.Arith.Linear.DenoteExpr
+import Lean.Meta.Tactic.Grind.Arith.Linear.Util
+import Lean.Meta.Tactic.Grind.Arith.Linear.IneqCnstr
+import Lean.Meta.Tactic.Grind.Arith.Linear.Proof
 public section
-
 namespace Lean.Meta.Grind.Arith.Linear
 
 def IneqCnstr.throwUnexpected (c : IneqCnstr) : LinearM α := do

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

@@ -4,12 +4,9 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Util
-
+public import Lean.Meta.Tactic.Grind.Arith.Linear.LinearM
 public section
-
 namespace Lean.Meta.Grind.Arith.Linear
 
 structure Case where

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

@@ -4,20 +4,19 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-public import Init.Grind.Ordered.Module
-public import Lean.Meta.Tactic.Grind.Simp
-public import Lean.Meta.Tactic.Grind.SynthInstance
-public import Lean.Meta.Tactic.Grind.Arith.Cutsat.ToInt
-public import Lean.Meta.Tactic.Grind.Arith.CommRing.RingId
-public import Lean.Meta.Tactic.Grind.Arith.CommRing.RingM
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Util
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Var
+public import Lean.Meta.Tactic.Grind.Types
+import Init.Grind.Ordered.Module
+import Lean.Meta.Tactic.Grind.Simp
+import Lean.Meta.Tactic.Grind.SynthInstance
+import Lean.Meta.Tactic.Grind.Arith.Cutsat.ToInt
+import Lean.Meta.Tactic.Grind.Arith.CommRing.RingId
+import Lean.Meta.Tactic.Grind.Arith.CommRing.RingM
+import Lean.Meta.Tactic.Grind.Arith.Linear.Util
+import Lean.Meta.Tactic.Grind.Arith.Linear.Var
 import Lean.Meta.Tactic.Grind.Arith.Insts
-
+import Init.Grind.Module.Envelope
 public section
-
 namespace Lean.Meta.Grind.Arith.Linear
 
 private def preprocess (e : Expr) : GoalM Expr := do
@@ -147,6 +146,54 @@ private def mkNoNatZeroDivInst? (u : Level) (type : Expr) : GoalM (Option Expr)
   let some natModuleInst ← synthInstance? natModuleType | return none
   synthInstance? <| mkApp2 (mkConst ``Grind.NoNatZeroDivisors [u]) type natModuleInst
 
+private def getInst? (declName : Name) (u : Level) (type : Expr) : GoalM (Option Expr) := do
+  synthInstance? <| mkApp (mkConst declName [u]) type
+
+private def getInst (declName : Name) (u : Level) (type : Expr) : GoalM Expr := do
+  synthInstance <| mkApp (mkConst declName [u]) type
+
+private def getBinHomoInst (declName : Name) (u : Level) (type : Expr) : GoalM Expr := do
+  synthInstance <| mkApp3 (mkConst declName [u, u, u]) type type type
+
+private def getHSMulIntInst (u : Level) (type : Expr) : GoalM Expr := do
+  synthInstance <| mkApp3 (mkConst ``HSMul [0, u, u]) Int.mkType type type
+
+private def getHSMulNatInst (u : Level) (type : Expr) : GoalM Expr := do
+  synthInstance <| mkApp3 (mkConst ``HSMul [0, u, u]) Nat.mkType type type
+
+private def checkToFieldDefEq? (leInst? parentInst? childInst? : Option Expr) (toFieldName : Name) (u : Level) (type : Expr) : GoalM (Option Expr) := do
+  let some leInst := leInst? | return none
+  let some parentInst := parentInst? | return none
+  let some childInst := childInst? | return none
+  let toField := mkApp3 (mkConst toFieldName [u]) type leInst childInst
+  unless (← isDefEqD parentInst toField) do
+    reportIssue! (← mkExpectedDefEqMsg parentInst toField)
+    return none
+  return some childInst
+
+private def ensureToFieldDefEq (parentInst : Expr) (inst : Expr) (toFieldName : Name) (u : Level) (type : Expr) : GoalM Unit := do
+  let toField := mkApp2 (mkConst toFieldName [u]) type inst
+  ensureDefEq parentInst toField
+
+private def ensureToHomoFieldDefEq (parentInst : Expr) (inst : Expr) (toFieldName : Name) (toHeteroName : Name) (u : Level) (type : Expr)
+    (extraType? : Option Expr := none) : GoalM Unit := do
+  let toField := mkApp2 (mkConst toFieldName [u]) type inst
+  let heteroToField :=
+    match extraType? with
+    | none => mkApp2 (mkConst toHeteroName [u]) type toField
+    | some extraType => mkApp3 (mkConst toHeteroName [0, u]) extraType type toField
+  ensureDefEq parentInst heteroToField
+
+private def getHSMulIntFn? (u : Level) (type : Expr) : GoalM (Option Expr) := do
+  let smulType := mkApp3 (mkConst ``HSMul [0, u, u]) Int.mkType type type
+  let some smulInst ← synthInstance? smulType | return none
+  internalizeFn <| mkApp4 (mkConst ``HSMul.hSMul [0, u, u]) Int.mkType type type smulInst
+
+private def getHSMulNatFn? (u : Level) (type : Expr) : GoalM (Option Expr) := do
+  let smulType := mkApp3 (mkConst ``HSMul [0, u, u]) Nat.mkType type type
+  let some smulInst ← synthInstance? smulType | return none
+  internalizeFn <| mkApp4 (mkConst ``HSMul.hSMul [0, u, u]) Nat.mkType type type smulInst
+
 def getStructId? (type : Expr) : GoalM (Option Nat) := do
   unless (← getConfig).linarith do return none
   if (← isCutsatType type) then return none
@@ -159,39 +206,9 @@ def getStructId? (type : Expr) : GoalM (Option Nat) := do
 where
   go? : GoalM (Option Nat) := do
     let u ← getDecLevel type
-    let rec getInst? (declName : Name) : GoalM (Option Expr) := do
-      synthInstance? <| mkApp (mkConst declName [u]) type
-    let rec getInst (declName : Name) : GoalM Expr := do
-      synthInstance <| mkApp (mkConst declName [u]) type
-    let rec getBinHomoInst (declName : Name) : GoalM Expr := do
-      synthInstance <| mkApp3 (mkConst declName [u, u, u]) type type type
-    let rec getHSMulIntInst : GoalM Expr := do
-      synthInstance <| mkApp3 (mkConst ``HSMul [0, u, u]) Int.mkType type type
-    let rec getHSMulNatInst : GoalM Expr := do
-      synthInstance <| mkApp3 (mkConst ``HSMul [0, u, u]) Nat.mkType type type
-    let rec checkToFieldDefEq? (leInst? parentInst? childInst? : Option Expr) (toFieldName : Name) : GoalM (Option Expr) := do
-      let some leInst := leInst? | return none
-      let some parentInst := parentInst? | return none
-      let some childInst := childInst? | return none
-      let toField := mkApp3 (mkConst toFieldName [u]) type leInst childInst
-      unless (← isDefEqD parentInst toField) do
-        reportIssue! (← mkExpectedDefEqMsg parentInst toField)
-        return none
-      return some childInst
-    let rec ensureToFieldDefEq (parentInst : Expr) (inst : Expr) (toFieldName : Name) : GoalM Unit := do
-      let toField := mkApp2 (mkConst toFieldName [u]) type inst
-      ensureDefEq parentInst toField
-    let rec ensureToHomoFieldDefEq (parentInst : Expr) (inst : Expr) (toFieldName : Name)
-        (toHeteroName : Name) (extraType? : Option Expr := none) : GoalM Unit := do
-      let toField := mkApp2 (mkConst toFieldName [u]) type inst
-      let heteroToField :=
-        match extraType? with
-        | none => mkApp2 (mkConst toHeteroName [u]) type toField
-        | some extraType => mkApp3 (mkConst toHeteroName [0, u]) extraType type toField
-      ensureDefEq parentInst heteroToField
     let ringId? ← CommRing.getRingId? type
-    let leInst? ← getInst? ``LE
-    let ltInst? ← getInst? ``LT
+    let leInst? ← getInst? ``LE u type
+    let ltInst? ← getInst? ``LT u type
     let lawfulOrderLTInst? ← mkLawfulOrderLTInst? u type ltInst? leInst?
     let isPreorderInst? ← mkIsPreorderInst? u type leInst?
     let isPartialInst? ← mkIsPartialOrderInst? u type leInst?
@@ -203,7 +220,7 @@ where
     let commRingInst? ← getCommRingInst? ringId?
     let ringInst? ← mkRingInst? u type commRingInst?
     let some intModuleInst ← mkIntModuleInst? u type ringInst? | return none
-    let addInst ← getBinHomoInst ``HAdd
+    let addInst ← getBinHomoInst ``HAdd u type
     let addFn ← internalizeFn <| mkApp4 (mkConst ``HAdd.hAdd [u, u, u]) type type type addInst
     let orderedAddInst? ← match leInst?, isPreorderInst? with
       | some leInst, some isPreorderInst =>
@@ -213,38 +230,38 @@ where
     -- preorderInst? may have been reset, check again whether this module is needed.
     if (← getConfig).ring && ringId?.isSome && isPreorderInst?.isNone then
       return none
-    let isPartialInst? ← checkToFieldDefEq? leInst? isPreorderInst? isPartialInst? ``Std.IsPartialOrder.toIsPreorder
-    let isLinearInst? ← checkToFieldDefEq? leInst? isPartialInst? isLinearInst? ``Std.IsLinearOrder.toIsPartialOrder
+    let isPartialInst? ← checkToFieldDefEq? leInst? isPreorderInst? isPartialInst? ``Std.IsPartialOrder.toIsPreorder u type
+    let isLinearInst? ← checkToFieldDefEq? leInst? isPartialInst? isLinearInst? ``Std.IsLinearOrder.toIsPartialOrder u type
     let addCommGroupInst := mkApp2 (mkConst ``Grind.IntModule.toAddCommGroup [u]) type intModuleInst
     let addCommMonoidInst := mkApp2 (mkConst ``Grind.AddCommGroup.toAddCommMonoid [u]) type addCommGroupInst
     let semiringInst? ← mkSemiringInst? u type ringInst?
-    let fieldInst? ← getInst? ``Grind.Field
+    let fieldInst? ← getInst? ``Grind.Field u type
     let one? ← mkOne? u type -- One must be created eagerly
     let orderedRingInst? ← mkOrderedRingInst? u type semiringInst? leInst? ltInst? isPreorderInst?
     let charInst? ← if let some semiringInst := semiringInst? then getIsCharInst? u type semiringInst else pure none
     let noNatDivInst? ← mkNoNatZeroDivInst? u type
     -- TODO: generate the remaining fields on demand
-    let zeroInst ← getInst ``Zero
+    let zeroInst ← getInst ``Zero u type
     let zero ← internalizeConst <| mkApp2 (mkConst ``Zero.zero [u]) type zeroInst
     let ofNatZeroType := mkApp2 (mkConst ``OfNat [u]) type (mkRawNatLit 0)
     let some ofNatZeroInst ← synthInstance? 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 subInst ← getBinHomoInst ``HSub
+    let subInst ← getBinHomoInst ``HSub u type
     let subFn ← internalizeFn <| mkApp4 (mkConst ``HSub.hSub [u, u, u]) type type type subInst
-    let negInst ← getInst ``Neg
+    let negInst ← getInst ``Neg u type
     let negFn ← internalizeFn <| mkApp2 (mkConst ``Neg.neg [u]) type negInst
-    let zsmulInst ← getHSMulIntInst
+    let zsmulInst ← getHSMulIntInst u type
     let zsmulFn ← internalizeFn <| mkApp4 (mkConst ``HSMul.hSMul [0, u, u]) Int.mkType type type zsmulInst
-    let nsmulInst ← getHSMulNatInst
+    let nsmulInst ← getHSMulNatInst u type
     let nsmulFn ← internalizeFn <| mkApp4 (mkConst ``HSMul.hSMul [0, u, u]) Nat.mkType type type nsmulInst
-    ensureToFieldDefEq zeroInst addCommMonoidInst ``Grind.AddCommMonoid.toZero
-    ensureToHomoFieldDefEq addInst addCommMonoidInst ``Grind.AddCommMonoid.toAdd ``instHAdd
-    ensureToHomoFieldDefEq subInst addCommGroupInst ``Grind.AddCommGroup.toSub ``instHSub
-    ensureToFieldDefEq negInst addCommGroupInst ``Grind.AddCommGroup.toNeg
-    ensureToHomoFieldDefEq zsmulInst intModuleInst ``Grind.IntModule.zsmul ``instHSMul (some Int.mkType)
-    ensureToHomoFieldDefEq nsmulInst intModuleInst ``Grind.IntModule.nsmul ``instHSMul (some Nat.mkType)
+    ensureToFieldDefEq zeroInst addCommMonoidInst ``Grind.AddCommMonoid.toZero u type
+    ensureToHomoFieldDefEq addInst addCommMonoidInst ``Grind.AddCommMonoid.toAdd ``instHAdd u type
+    ensureToHomoFieldDefEq subInst addCommGroupInst ``Grind.AddCommGroup.toSub ``instHSub u type
+    ensureToFieldDefEq negInst addCommGroupInst ``Grind.AddCommGroup.toNeg u type
+    ensureToHomoFieldDefEq zsmulInst intModuleInst ``Grind.IntModule.zsmul ``instHSMul u type (some Int.mkType)
+    ensureToHomoFieldDefEq nsmulInst intModuleInst ``Grind.IntModule.nsmul ``instHSMul u type (some Nat.mkType)
     let leFn? ← if let some leInst := leInst? then
       some <$> (internalizeFn <| mkApp2 (mkConst ``LE.le [u]) type leInst)
     else
@@ -253,24 +270,17 @@ where
       some <$> (internalizeFn <| mkApp2 (mkConst ``LT.lt [u]) type ltInst)
     else
       pure none
-    let rec getHSMulIntFn? : GoalM (Option Expr) := do
-      let smulType := mkApp3 (mkConst ``HSMul [0, u, u]) Int.mkType type type
-      let some smulInst ← synthInstance? smulType | return none
-      internalizeFn <| mkApp4 (mkConst ``HSMul.hSMul [0, u, u]) Int.mkType type type smulInst
-    let rec getHSMulNatFn? : GoalM (Option Expr) := do
-      let smulType := mkApp3 (mkConst ``HSMul [0, u, u]) Nat.mkType type type
-      let some smulInst ← synthInstance? smulType | return none
-      internalizeFn <| mkApp4 (mkConst ``HSMul.hSMul [0, u, u]) Nat.mkType type type smulInst
-    let zsmulFn? ← getHSMulIntFn?
-    let nsmulFn? ← getHSMulNatFn?
+    let zsmulFn? ← getHSMulIntFn? u type
+    let nsmulFn? ← getHSMulNatFn? u type
     let homomulFn? ← if commRingInst?.isSome then
-      let mulInst ← getBinHomoInst ``HMul
+      let mulInst ← getBinHomoInst ``HMul u type
       pure <| some (← internalizeFn <| mkApp4 (mkConst ``HMul.hMul [u, u, u]) type type type mulInst)
     else
       pure none
     let id := (← get').structs.size
     let struct : Struct := {
-      id, type, u, intModuleInst, leInst?, ltInst?, lawfulOrderLTInst?, isPreorderInst?, orderedAddInst?, isPartialInst?, isLinearInst?, noNatDivInst?
+      id, type, u, intModuleInst, leInst?, ltInst?, lawfulOrderLTInst?, isPreorderInst?,
+      orderedAddInst?, isPartialInst?, isLinearInst?, noNatDivInst?
       leFn?, ltFn?, addFn, subFn, negFn, zsmulFn, nsmulFn, zsmulFn?, nsmulFn?, zero, one?
       ringInst?, commRingInst?, orderedRingInst?, charInst?, ringId?, fieldInst?, ofNatZero, homomulFn?
     }
@@ -286,7 +296,52 @@ where
           -- Create `1` variable, and assert `0 ≠ 1`
           let x ← mkVar one (mark := false)
           addZeroNeOne x
+    return some id
 
+private def mkNatModuleInst? (u : Level) (type : Expr) : GoalM (Option Expr) := do
+  synthInstance? <| mkApp (mkConst ``Grind.NatModule [u]) type
+
+def getNatStructId? (type : Expr) : GoalM (Option Nat) := do
+  unless (← getConfig).linarith do return none
+  if (← isCutsatType type) then return none
+  if let some id? := (← get').natTypeIdOf.find? { expr := type } then
+    return id?
+  else
+    let id? ← go?
+    modify' fun s => { s with natTypeIdOf := s.natTypeIdOf.insert { expr := type } id? }
+    return id?
+where
+  go? : GoalM (Option Nat) := do
+    let u ← getDecLevel type
+    let some natModuleInst ← mkNatModuleInst? u type | return none
+    let q ← shareCommon (← canon (mkApp2 (mkConst ``Grind.IntModule.OfNatModule.Q [u]) type natModuleInst))
+    let some structId ← getStructId? q
+      | throwError "`grind` unexpected failure, failure to initialize auxiliary `IntModule`{indentExpr q}"
+    let leInst? ← getInst? ``LE u type
+    let ltInst? ← getInst? ``LT u type
+    let isPreorderInst? ← mkIsPreorderInst? u type leInst?
+    let lawfulOrderLTInst? ← mkLawfulOrderLTInst? u type ltInst? leInst?
+    let addInst ← getBinHomoInst ``HAdd u type
+    let addFn ← internalizeFn <| mkApp4 (mkConst ``HAdd.hAdd [u, u, u]) type type type addInst
+    let orderedAddInst? ← match leInst?, isPreorderInst? with
+      | some leInst, some isPreorderInst =>
+        synthInstance? <| mkApp4 (mkConst ``Grind.OrderedAdd [u]) type addInst leInst isPreorderInst
+      | _, _ => pure none
+    let addInst' ← synthInstance <| mkApp (mkConst ``Add [u]) type
+    let addRightCancelInst? ← synthInstance? <| mkApp2 (mkConst ``Grind.AddRightCancel [u]) type addInst'
+    let toQFn ← internalizeFn <| mkApp2 (mkConst ``Grind.IntModule.OfNatModule.toQ [u]) type natModuleInst
+    let zeroInst ← getInst ``Zero u type
+    let zero ← internalizeConst <| mkApp2 (mkConst ``Zero.zero [u]) type zeroInst
+    let smulInst ← getHSMulNatInst u type
+    let smulFn ← internalizeFn <| mkApp4 (mkConst ``HSMul.hSMul [0, u, u]) Nat.mkType type type smulInst
+    let rfl_q := mkApp (mkConst ``Eq.refl [.succ u]) q
+    let id := (← get').natStructs.size
+    let natStruct : NatStruct := {
+      id, structId, u, type, natModuleInst,
+      leInst?, ltInst?, lawfulOrderLTInst?, isPreorderInst?, orderedAddInst?, addRightCancelInst?,
+      rfl_q, zero, toQFn, addFn, smulFn
+    }
+    modify' fun s => { s with natStructs := s.natStructs.push natStruct }
     return some id
 
 end Lean.Meta.Grind.Arith.Linear

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

@@ -4,16 +4,12 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Init.Grind.Ordered.Linarith
 public import Lean.ToExpr
-
 public section
-
 namespace Lean.Meta.Grind.Arith.Linear
 open Grind.Linarith
-
 /-!
 `ToExpr` instances for `Linarith.Poly` types.
 -/

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

@@ -4,16 +4,13 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Init.Grind.Ring.Poly
 public import Init.Grind.Ordered.Linarith
 public import Lean.Data.PersistentArray
 public import Lean.Meta.Tactic.Grind.ExprPtr
 public import Init.Data.Rat.Basic
-
 public section
-
 namespace Lean.Meta.Grind.Arith.Linear
 export Lean.Grind.Linarith (Var Poly)
 
@@ -64,6 +61,7 @@ structure DiseqCnstr where
 inductive DiseqCnstrProof where
   | core (a b : Expr) (lhs rhs : LinExpr)
   | coreCommRing (a b : Expr) (ra rb : Grind.CommRing.Expr) (p : Grind.CommRing.Poly) (lhs' : LinExpr)
+  | coreOfNat (a b : Expr) (natStructId : Nat) (lhs rhs : LinExpr)
   | neg (c : DiseqCnstr)
   | subst (k₁ k₂ : Int) (c₁ : EqCnstr) (c₂ : DiseqCnstr)
   | subst1 (k : Int) (c₁ : EqCnstr) (c₂ : DiseqCnstr)
@@ -89,53 +87,53 @@ Each type must at least implement the instance `IntModule`.
 For being able to process inequalities, it must at least implement `Preorder`, and `OrderedAdd`
 -/
 structure Struct where
-  id               : Nat
+  id                 : Nat
   /-- If the structure is a ring, we store its id in the `CommRing` module at `ringId?` -/
-  ringId?          : Option Nat
-  type             : Expr
+  ringId?            : Option Nat
+  type               : Expr
   /-- Cached `getDecLevel type` -/
-  u                : Level
+  u                  : Level
   /-- `IntModule` instance -/
-  intModuleInst    : Expr
+  intModuleInst      : Expr
   /-- `LE` instance if available -/
-  leInst?          : Option Expr
+  leInst?            : Option Expr
   /-- `LT` instance if available -/
-  ltInst?          : Option Expr
+  ltInst?            : Option Expr
   /-- `LawfulOrderLT` instance if available -/
-  lawfulOrderLTInst?    : Option Expr
+  lawfulOrderLTInst? : Option Expr
   /-- `IsPreorder` instance if available -/
   isPreorderInst?    : Option Expr
   /-- `OrderedAdd` instance with `IsPreorder` if available -/
-  orderedAddInst?       : Option Expr
+  orderedAddInst?    : Option Expr
   /-- `IsPartialOrder` instance if available -/
   isPartialInst?     : Option Expr
   /-- `IsLinearOrder` instance if available -/
   isLinearInst?      : Option Expr
   /-- `NoNatZeroDivisors` -/
-  noNatDivInst?    : Option Expr
+  noNatDivInst?      : Option Expr
   /-- `Ring` instance -/
-  ringInst?        : Option Expr
+  ringInst?          : Option Expr
   /-- `CommRing` instance -/
-  commRingInst?    : Option Expr
+  commRingInst?      : Option Expr
   /-- `OrderedRing` instance with `Preorder` -/
   orderedRingInst?   : Option Expr
   /-- `Field` instance -/
-  fieldInst?       : Option Expr
+  fieldInst?         : Option Expr
   /-- `IsCharP` instance for `type` if available. -/
-  charInst?        : Option (Expr × Nat)
-  zero             : Expr
-  ofNatZero        : Expr
-  one?             : Option Expr
-  leFn?            : Option Expr
-  ltFn?            : Option Expr
-  addFn            : Expr
-  zsmulFn          : Expr
-  nsmulFn          : Expr
-  zsmulFn?         : Option Expr
-  nsmulFn?         : Option Expr
-  homomulFn?       : Option Expr -- homogeneous multiplication if structure is a ring
-  subFn            : Expr
-  negFn            : Expr
+  charInst?          : Option (Expr × Nat)
+  zero               : Expr
+  ofNatZero          : Expr
+  one?               : Option Expr
+  leFn?              : Option Expr
+  ltFn?              : Option Expr
+  addFn              : Expr
+  zsmulFn            : Expr
+  nsmulFn            : Expr
+  zsmulFn?           : Option Expr
+  nsmulFn?           : Option Expr
+  homomulFn?         : Option Expr -- homogeneous multiplication if structure is a ring
+  subFn              : Expr
+  negFn              : Expr
   /--
   Mapping from variables to their denotations.
   Remark each variable can be in only one ring.
@@ -203,6 +201,33 @@ structure Struct where
   ignored : PArray Expr := {}
   deriving Inhabited
 
+structure NatStruct where
+  id                  : Nat
+  /-- Id for `OfNatModule.Q` -/
+  structId            : Nat
+  type                : Expr
+  /-- Cached `getDecLevel type` -/
+  u                   : Level
+  /-- `NatModule` instance for `type` -/
+  natModuleInst       : Expr
+  /-- `LE` instance if available -/
+  leInst?             : Option Expr
+  /-- `LT` instance if available -/
+  ltInst?             : Option Expr
+  /-- `LawfulOrderLT` instance if available -/
+  lawfulOrderLTInst?  : Option Expr
+  /-- `IsPreorder` instance if available -/
+  isPreorderInst?     : Option Expr
+  /-- `OrderedAdd` instance with `IsPreorder` if available -/
+  orderedAddInst?     : Option Expr
+  addRightCancelInst? : Option Expr
+  rfl_q               : Expr -- `@Eq.Refl (OfNatModule.Q type)`
+  zero                : Expr
+  toQFn               : Expr
+  addFn               : Expr
+  smulFn              : Expr
+  termMap             : PHashMap ExprPtr (Expr × Expr) := {}
+
 /-- State for all `IntModule` types detected by `grind`. -/
 structure State where
   /--
@@ -216,6 +241,16 @@ structure State where
   typeIdOf : PHashMap ExprPtr (Option Nat) := {}
   /- Mapping from expressions/terms to their structure ids. -/
   exprToStructId : PHashMap ExprPtr Nat := {}
+  /-- `NatModule`. We support them using the envelope `OfNatModule.Q` -/
+  natStructs : Array NatStruct := {}
+  /--
+  Mapping from types to its "nat module id". We cache failures using `none`.
+  `natTypeIdOf[type]` is `some id`, then `id < natStructs.size`.
+  If a type is in this map, it is not in `typeIdOf`.
+  -/
+  natTypeIdOf : PHashMap ExprPtr (Option Nat) := {}
+  /- Mapping from expressions/terms to their nat structure ids. -/
+  exprToNatStructId : PHashMap ExprPtr Nat := {}
   deriving Inhabited
 
 end Lean.Meta.Grind.Arith.Linear

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

@@ -4,81 +4,11 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-public import Lean.Meta.Tactic.Grind.Types
-public import Lean.Meta.Tactic.Grind.Arith.CommRing.RingM
-
+public import Lean.Meta.Tactic.Grind.Arith.Linear.LinearM
 public section
-
 namespace Lean.Meta.Grind.Arith.Linear
 
-def get' : GoalM State := do
-  return (← get).arith.linear
-
-@[inline] def modify' (f : State → State) : GoalM Unit := do
-  modify fun s => { s with arith.linear := f s.arith.linear }
-
-structure LinearM.Context where
-  structId : Nat
-
-class MonadGetStruct (m : Type → Type) where
-  getStruct : m Struct
-
-export MonadGetStruct (getStruct)
-
-@[always_inline]
-instance (m n) [MonadLift m n] [MonadGetStruct m] : MonadGetStruct n where
-  getStruct    := liftM (getStruct : m Struct)
-
-/-- We don't want to keep carrying the `StructId` around. -/
-abbrev LinearM := ReaderT LinearM.Context GoalM
-
-abbrev LinearM.run (structId : Nat) (x : LinearM α) : GoalM α :=
-  x { structId }
-
-abbrev getStructId : LinearM Nat :=
-  return (← read).structId
-
-protected def LinearM.getStruct : LinearM Struct := do
-  let s ← get'
-  let structId ← getStructId
-  if h : structId < s.structs.size then
-    return s.structs[structId]
-  else
-    throwError "`grind` internal error, invalid structure id"
-
-instance : MonadGetStruct LinearM where
-  getStruct := LinearM.getStruct
-
-open CommRing
-
-def getRingCore? (ringId? : Option Nat) : GoalM (Option Ring) := do
-  let some ringId := ringId? | return none
-  RingM.run ringId do return some (← getRing)
-
-def throwNotRing : LinearM α :=
-  throwError "`grind linarith` internal error, structure is not a ring"
-
-def throwNotCommRing : LinearM α :=
-  throwError "`grind linarith` internal error, structure is not a commutative ring"
-
-def getRing? : LinearM (Option Ring) := do
-  getRingCore? (← getStruct).ringId?
-
-def getRing : LinearM Ring := do
-  let some ring ← getRing?
-    | throwNotCommRing
-  return ring
-
-instance : MonadRing LinearM where
-  getRing := Linear.getRing
-  modifyRing f := do
-    let some ringId := (← getStruct).ringId? | throwNotCommRing
-    RingM.run ringId do modifyRing f
-  canonExpr e := do shareCommon (← canon e)
-  synthInstance? e := Grind.synthInstance? e
-
 def getZero : LinearM Expr :=
   return (← getStruct).zero
 
@@ -87,11 +17,6 @@ def getOne : LinearM Expr := do
     | throwNotRing
   return one
 
-def withRingM (x : RingM α) : LinearM α := do
-  let some ringId := (← getStruct).ringId?
-    | throwNotCommRing
-  RingM.run ringId x
-
 def isCommRing : LinearM Bool :=
   return (← getStruct).ringId?.isSome
 
@@ -104,10 +29,6 @@ def isLinearOrder : LinearM Bool :=
 def hasNoNatZeroDivisors : LinearM Bool :=
   return (← getStruct).noNatDivInst?.isSome
 
-@[inline] def modifyStruct (f : Struct → Struct) : LinearM Unit := do
-  let structId ← getStructId
-  modify' fun s => { s with structs := s.structs.modify structId f }
-
 def getTermStructId? (e : Expr) : GoalM (Option Nat) := do
   return (← get').exprToStructId.find? { expr := e }
 

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

@@ -4,12 +4,10 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
-public import Lean.Meta.Tactic.Grind.Arith.Linear.Util
-
+public import Lean.Meta.Tactic.Grind.Arith.Linear.LinearM
+import Lean.Meta.Tactic.Grind.Arith.Linear.Util
 public section
-
 namespace Lean.Meta.Grind.Arith.Linear
 
 def mkVar (e : Expr) (mark := true) : LinearM Var := do

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

@@ -4,11 +4,9 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 module
-
 prelude
 public import Init.Grind.Ordered.Linarith
 public import Lean.Meta.Tactic.Grind.VarRename
-
 namespace Lean.Grind.Linarith
 open Lean.Meta.Grind
 

tests/lean/run/grind_nat_module.lean

@@ -0,0 +1,5 @@
+open Lean Grind
+variable (M : Type) [NatModule M] [AddRightCancel M]
+
+example (x y : M) : 2 • x + 3 • y + x = 3 • (x + y) := by
+  grind