From 5daf65ec56b7d41fa809138550e338802e2a5bd7 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Fri, 22 Aug 2025 13:36:05 -0700
Subject: [PATCH] feat: add helper theorems for `NatModule` (#10069)

This PR adds helper theorems to support `NatModule` in `grind linarith`.
---
 src/Init/Grind/Module.lean             |  1 +
 src/Init/Grind/Module/Envelope.lean    |  6 ++++
 src/Init/Grind/Module/OfNatModule.lean | 49 ++++++++++++++++++++++++++
 3 files changed, 56 insertions(+)
 create mode 100644 src/Init/Grind/Module/OfNatModule.lean

diff --git a/src/Init/Grind/Module.lean b/src/Init/Grind/Module.lean
index a314e49967..7116319912 100644
--- a/src/Init/Grind/Module.lean
+++ b/src/Init/Grind/Module.lean
@@ -8,5 +8,6 @@ module
 prelude
 public import Init.Grind.Module.Basic
 public import Init.Grind.Module.Envelope
+public import Init.Grind.Module.OfNatModule
 
 public section
diff --git a/src/Init/Grind/Module/Envelope.lean b/src/Init/Grind/Module/Envelope.lean
index 6710887268..17d12cba8a 100644
--- a/src/Init/Grind/Module/Envelope.lean
+++ b/src/Init/Grind/Module/Envelope.lean
@@ -222,6 +222,12 @@ attribute [instance] ofNatModule
 theorem toQ_add (a b : α) : toQ (a + b) = toQ a + toQ b := by
   simp; apply Quot.sound; simp
 
+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
+  simp; apply Quot.sound; simp; exists 0
+
 /-!
 Helper definitions and theorems for proving `toQ` is injective when
 `CommSemiring` has the right_cancel property
diff --git a/src/Init/Grind/Module/OfNatModule.lean b/src/Init/Grind/Module/OfNatModule.lean
new file mode 100644
index 0000000000..9d981085e7
--- /dev/null
+++ b/src/Init/Grind/Module/OfNatModule.lean
@@ -0,0 +1,49 @@
+/-
+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
+import Init.Grind.Module.Envelope
+
+namespace Lean.Grind.IntModule.OfNatModule
+
+/-!
+Support for `NatModule` in the `grind` linear arithmetic module.
+-/
+
+theorem of_eq {α} [NatModule α] {a b : α} {a' b' : Q α}
+    (h₁ : toQ a = a') (h₂ : toQ b = b') : a = b → a' = b' := by
+  intro h; rw [← h₁, ← h₂, h]
+
+theorem of_diseq {α} [NatModule α] [AddRightCancel α] {a b : α} {a' b' : Q α}
+    (h₁ : toQ a = a') (h₂ : toQ b = b') : a ≠ b → a' ≠ b' := by
+  rw [← h₁, ← h₂]; intro h₃ h₄; replace h₄ := toQ_inj h₄; contradiction
+
+theorem of_le {α} [NatModule α] [LE α] [LT α] [Preorder α] [OrderedAdd α] {a b : α} {a' b' : Q α}
+    (h₁ : toQ a = a') (h₂ : toQ b = b') : a ≤ b → a' ≤ b' := by
+  rw [← h₁, ← h₂, toQ_le]; intro; assumption
+
+theorem of_not_le {α} [NatModule α] [LE α] [LT α] [Preorder α] [OrderedAdd α] {a b : α} {a' b' : Q α}
+    (h₁ : toQ a = a') (h₂ : toQ b = b') : ¬ a ≤ b → ¬ a' ≤ b' := by
+  rw [← h₁, ← h₂, toQ_le]; intro; assumption
+
+theorem of_lt {α} [NatModule α] [LE α] [LT α] [Preorder α] [OrderedAdd α] {a b : α} {a' b' : Q α}
+    (h₁ : toQ a = a') (h₂ : toQ b = b') : a < b → a' < b' := by
+  rw [← h₁, ← h₂, toQ_lt]; intro; assumption
+
+theorem of_not_lt {α} [NatModule α] [LE α] [LT α] [Preorder α] [OrderedAdd α] {a b : α} {a' b' : Q α}
+    (h₁ : toQ a = a') (h₂ : toQ b = b') : ¬ a < b → ¬ a' < b' := by
+  rw [← h₁, ← h₂, toQ_lt]; intro; assumption
+
+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]
+
+end Lean.Grind.IntModule.OfNatModule