From 64926ba2e9c7768718683a048c3e5d7e78b0301d Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Thu, 4 Dec 2025 21:24:30 -0800
Subject: [PATCH] feat: `grind` propagators for `Nat` operators
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This PR implements `grind` propagators for `Nat` operators that have a
simproc associated with them, but do not have any theory solver
support. Examples:

```lean
example (a b : Nat) : a = 3 → b = 6 → a &&& b = 2 := by grind
example (a b : Nat) : a = 3 → b = 6 → a ||| b = 7 := by grind
example (a b : Nat) : a = 3 → b = 6 → a ^^^ b = 5 := by grind
example (a b : Nat) : a = 3 → b = 6 → a <<< b = 192 := by grind
example (a b : Nat) : a = 1135 → b = 6 → a >>> b = 17 := by grind
```

Closes #11498
---
 src/Init/Grind/Lemmas.lean                    |  8 +++
 src/Lean/Meta/Tactic/Grind/Arith.lean         |  1 +
 .../Meta/Tactic/Grind/Arith/Propagate.lean    | 65 +++++++++++++++++++
 tests/lean/run/grind_11498.lean               | 16 +++++
 4 files changed, 90 insertions(+)
 create mode 100644 src/Lean/Meta/Tactic/Grind/Arith/Propagate.lean
 create mode 100644 tests/lean/run/grind_11498.lean

diff --git a/src/Init/Grind/Lemmas.lean b/src/Init/Grind/Lemmas.lean
index 8acf3f39568a..1d609c247588 100644
--- a/src/Init/Grind/Lemmas.lean
+++ b/src/Init/Grind/Lemmas.lean
@@ -185,4 +185,12 @@ theorem decide_eq_false {p : Prop} {_ : Decidable p} : p = False → decide p =
 theorem of_lookahead (p : Prop) (h : (¬ p) → False) : p = True := by
   simp at h; simp [h]
 
+/-! Nat propagators -/
+
+theorem Nat.and_congr {a b : Nat} {k₁ k₂ k : Nat} (h₁ : a = k₁) (h₂ : b = k₂) : k == k₁ &&& k₂ → a &&& b = k := by simp_all
+theorem Nat.xor_congr {a b : Nat} {k₁ k₂ k : Nat} (h₁ : a = k₁) (h₂ : b = k₂) : k == k₁ ^^^ k₂ → a ^^^ b = k := by simp_all
+theorem Nat.or_congr {a b : Nat} {k₁ k₂ k : Nat} (h₁ : a = k₁) (h₂ : b = k₂) : k == k₁ ||| k₂ → a ||| b = k := by simp_all
+theorem Nat.shiftLeft_congr {a b : Nat} {k₁ k₂ k : Nat} (h₁ : a = k₁) (h₂ : b = k₂) : k == k₁ <<< k₂ → a <<< b = k := by simp_all
+theorem Nat.shiftRight_congr {a b : Nat} {k₁ k₂ k : Nat} (h₁ : a = k₁) (h₂ : b = k₂) : k == k₁ >>> k₂ → a >>> b = k := by simp_all
+
 end Lean.Grind
diff --git a/src/Lean/Meta/Tactic/Grind/Arith.lean b/src/Lean/Meta/Tactic/Grind/Arith.lean
index 1431069c2518..5c916589993a 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith.lean
@@ -12,3 +12,4 @@ public import Lean.Meta.Tactic.Grind.Arith.CommRing
 public import Lean.Meta.Tactic.Grind.Arith.Linear
 public import Lean.Meta.Tactic.Grind.Arith.Cutsat
 public import Lean.Meta.Tactic.Grind.Arith.Simproc
+public import Lean.Meta.Tactic.Grind.Arith.Propagate
diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Propagate.lean b/src/Lean/Meta/Tactic/Grind/Arith/Propagate.lean
new file mode 100644
index 000000000000..dc27dcfc8afe
--- /dev/null
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Propagate.lean
@@ -0,0 +1,65 @@
+/-
+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
+import Init.Grind
+import Lean.Meta.Tactic.Grind.PropagatorAttr
+public section
+namespace Lean.Meta.Grind
+
+/-!
+This file defines propagators for `Nat` operators that have simprocs associated with them, but do not
+have support in satellite solvers. The goal is to workaround a nasty interaction between
+E-matching and normalization. Consider the following example:
+```
+@[grind] def mask := 15
+
+@[grind =] theorem foo (x : Nat) : x &&& 15 = x % 16 := by sorry
+
+example (x : Nat) : 3 &&& mask = 3 := by
+  grind only [foo, mask]
+```
+This is a very simple version of issue #11498.
+The e-graph contains `3 &&& mask`. E-matching finds that `3 &&& mask` matches pattern `x &&& 15`
+modulo the equality `mask = 15`, binding `x := 3`.
+Thus is produces the theorem instance `3 &&& 15 = 3 % 16`, which simplifies to `True` using the
+builtin simprocs for `&&&` and `%`. So, nothing is learned.
+**Tension**: The invariant "all terms in e-graph are normalized" conflicts with congruence needing
+the intermediate term `3 &&& 15` to make the connection.
+This is yet another example that shows how tricky is to combine e-matching with a normalizer.
+It is yet another reason for not letting users to extend the normalizer.
+If we do, we should be able to mark which ones should be used not normalize theorem instances.
+The following propagator ensure that `3 &&& mask` is merged with the equivalence class `3` if
+`mask = 15`.
+-/
+
+def propagateNatBinOp (declName : Name) (congrThmName : Name) (op : Nat → Nat → Nat) (e : Expr) : GoalM Unit := do
+  let arity := 6
+  unless e.isAppOfArity declName arity do return ()
+  unless e.getArg! 0 |>.isConstOf ``Nat do return ()
+  unless e.getArg! 1 |>.isConstOf ``Nat do return ()
+  let a := e.getArg! (arity - 2) arity
+  let aRoot ← getRoot a
+  let some k₁ ← getNatValue? aRoot | return ()
+  let b := e.getArg! (arity - 1) arity
+  let bRoot ← getRoot b
+  let some k₂ ← getNatValue? bRoot | return ()
+  let k := op k₁ k₂
+  let r ← shareCommon (mkNatLit k)
+  internalize r 0
+  let h := mkApp8 (mkConst congrThmName) a b aRoot bRoot r (← mkEqProof a aRoot) (← mkEqProof b bRoot) eagerReflBoolTrue
+  pushEq e r h
+
+builtin_grind_propagator propagateNatAnd ↑HAnd.hAnd := propagateNatBinOp ``HAnd.hAnd ``Grind.Nat.and_congr (· &&& ·)
+builtin_grind_propagator propagateNatOr ↑HOr.hOr := propagateNatBinOp ``HOr.hOr ``Grind.Nat.or_congr (· ||| ·)
+builtin_grind_propagator propagateNatXOr ↑HXor.hXor := propagateNatBinOp ``HXor.hXor ``Grind.Nat.xor_congr (· ^^^ ·)
+builtin_grind_propagator propagateNatShiftLeft ↑HShiftLeft.hShiftLeft :=
+  propagateNatBinOp ``HShiftLeft.hShiftLeft ``Grind.Nat.shiftLeft_congr (· <<< ·)
+builtin_grind_propagator propagateNatShiftRight ↑HShiftRight.hShiftRight :=
+  propagateNatBinOp ``HShiftRight.hShiftRight ``Grind.Nat.shiftRight_congr (· >>> ·)
+
+end Lean.Meta.Grind
diff --git a/tests/lean/run/grind_11498.lean b/tests/lean/run/grind_11498.lean
new file mode 100644
index 000000000000..59f259a175e2
--- /dev/null
+++ b/tests/lean/run/grind_11498.lean
@@ -0,0 +1,16 @@
+@[grind]
+def RMASK : Nat := 65535
+
+@[grind =] theorem and_65535_eq_mod (x : Nat) : x &&& 65535 = x % 65536 := by sorry
+
+example : 3327 = 3327 &&& RMASK := by
+  grind
+
+example : 3327 = 3327 &&& RMASK := by
+  grind only [RMASK] -- and_... is not needed
+
+example (a b : Nat) : a = 3 → b = 6 → a &&& b = 2 := by grind
+example (a b : Nat) : a = 3 → b = 6 → a ||| b = 7 := by grind
+example (a b : Nat) : a = 3 → b = 6 → a ^^^ b = 5 := by grind
+example (a b : Nat) : a = 3 → b = 6 → a <<< b = 192 := by grind
+example (a b : Nat) : a = 1135 → b = 6 → a >>> b = 17 := by grind