From 535ce0b8fdf545a6e0c62ff3fd1634d4b07d32ef Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Tue, 1 Jul 2025 10:52:16 -0700
Subject: [PATCH] feat: `pow_add` normalization in `grind` (#9133)

This PR adds support for `a^(m+n)` in the `grind` normalizer.
---
 src/Lean/Expr.lean                            |  7 +++
 src/Lean/Meta/Tactic/Grind/Arith.lean         |  1 +
 src/Lean/Meta/Tactic/Grind/Arith/Simproc.lean | 23 ++++++++++
 src/Lean/Meta/Tactic/Grind/SimpUtil.lean      |  2 +
 tests/lean/grind/algebra/hyperoperations.lean | 45 -------------------
 tests/lean/run/grind_bitvec2.lean             |  2 +-
 tests/lean/run/grind_hyper_ex.lean            |  3 +-
 7 files changed, 35 insertions(+), 48 deletions(-)
 create mode 100644 src/Lean/Meta/Tactic/Grind/Arith/Simproc.lean
 delete mode 100644 tests/lean/grind/algebra/hyperoperations.lean

diff --git a/src/Lean/Expr.lean b/src/Lean/Expr.lean
index 7a3202b169..59012bed65 100644
--- a/src/Lean/Expr.lean
+++ b/src/Lean/Expr.lean
@@ -2198,6 +2198,9 @@ private def natSubFn : Expr :=
 private def natMulFn : Expr :=
   mkApp4 (mkConst ``HMul.hMul [0, 0, 0]) Nat.mkType Nat.mkType Nat.mkType Nat.mkInstHMul
 
+private def natPowFn : Expr :=
+  mkApp4 (mkConst ``HPow.hPow [0, 0, 0]) Nat.mkType Nat.mkType Nat.mkType Nat.mkInstHPow
+
 /-- Given `a : Nat`, returns `Nat.succ a` -/
 def mkNatSucc (a : Expr) : Expr :=
   mkApp (mkConst ``Nat.succ) a
@@ -2214,6 +2217,10 @@ def mkNatSub (a b : Expr) : Expr :=
 def mkNatMul (a b : Expr) : Expr :=
   mkApp2 natMulFn a b
 
+/-- Given `a b : Nat`, returns `a ^ b` -/
+def mkNatPow (a b : Expr) : Expr :=
+  mkApp2 natPowFn a b
+
 private def natLEPred : Expr :=
   mkApp2 (mkConst ``LE.le [0]) Nat.mkType Nat.mkInstLE
 
diff --git a/src/Lean/Meta/Tactic/Grind/Arith.lean b/src/Lean/Meta/Tactic/Grind/Arith.lean
index bc2f6fe3f6..53ddb17309 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith.lean
@@ -12,3 +12,4 @@ import Lean.Meta.Tactic.Grind.Arith.Offset
 import Lean.Meta.Tactic.Grind.Arith.Cutsat
 import Lean.Meta.Tactic.Grind.Arith.CommRing
 import Lean.Meta.Tactic.Grind.Arith.Linear
+import Lean.Meta.Tactic.Grind.Arith.Simproc
diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Simproc.lean b/src/Lean/Meta/Tactic/Grind/Arith/Simproc.lean
new file mode 100644
index 0000000000..d3f7ee0996
--- /dev/null
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Simproc.lean
@@ -0,0 +1,23 @@
+/-
+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
+-/
+prelude
+import Init.Simproc
+import Lean.Meta.Tactic.Simp.Simproc
+
+namespace Lean.Meta.Grind.Arith
+
+/-- Applies `a^(m+n) = a^m * a^n` -/
+builtin_simproc_decl reducePowAdd ((_ ^ _ : Nat)) := fun e => do
+  let_expr HPow.hPow _ _ _ _ a k := e | return .continue
+  let_expr HAdd.hAdd _ _ _ _ m n := k | return .continue
+  let r := mkNatMul (mkNatPow a m) (mkNatPow a n)
+  let h := mkApp3 (mkConst ``Nat.pow_add) a m n
+  return .visit { expr := r, proof? := some h }
+
+def addSimproc (s : Simprocs) : CoreM Simprocs := do
+  s.add ``reducePowAdd (post := true)
+
+end Lean.Meta.Grind.Arith
diff --git a/src/Lean/Meta/Tactic/Grind/SimpUtil.lean b/src/Lean/Meta/Tactic/Grind/SimpUtil.lean
index 247f9fb1a0..5c8114dff4 100644
--- a/src/Lean/Meta/Tactic/Grind/SimpUtil.lean
+++ b/src/Lean/Meta/Tactic/Grind/SimpUtil.lean
@@ -8,6 +8,7 @@ import Lean.Meta.Tactic.Simp.Simproc
 import Lean.Meta.Tactic.Grind.Simp
 import Lean.Meta.Tactic.Grind.MatchDiscrOnly
 import Lean.Meta.Tactic.Grind.MatchCond
+import Lean.Meta.Tactic.Grind.Arith.Simproc
 import Lean.Meta.Tactic.Simp.BuiltinSimprocs.List
 
 namespace Lean.Meta.Grind
@@ -66,6 +67,7 @@ protected def getSimprocs : MetaM (Array Simprocs) := do
   let s := s.erase ``List.reduceReplicate
   let s ← addSimpMatchDiscrsOnly s
   let s ← addPreMatchCondSimproc s
+  let s ← Arith.addSimproc s
   let s ← s.add ``simpBoolEq (post := false)
   return #[s]
 
diff --git a/tests/lean/grind/algebra/hyperoperations.lean b/tests/lean/grind/algebra/hyperoperations.lean
deleted file mode 100644
index 35df5965e7..0000000000
--- a/tests/lean/grind/algebra/hyperoperations.lean
+++ /dev/null
@@ -1,45 +0,0 @@
-abbrev ℕ := Nat
-
-def hyperoperation : ℕ → ℕ → ℕ → ℕ
-  | 0, _, k => k + 1
-  | 1, m, 0 => m
-  | 2, _, 0 => 0
-  | _ + 3, _, 0 => 1
-  | n + 1, m, k + 1 => hyperoperation n m (hyperoperation (n + 1) m k)
-
-attribute [local grind] hyperoperation
-
-@[grind =]
-theorem hyperoperation_zero (m k : ℕ) : hyperoperation 0 m k = k + 1 := by
-  grind
-
-@[grind =]
-theorem hyperoperation_recursion (n m k : ℕ) :
-    hyperoperation (n + 1) m (k + 1) = hyperoperation n m (hyperoperation (n + 1) m k) := by
-  grind
-
-@[grind =]
-theorem hyperoperation_one (m k : ℕ) : hyperoperation 1 m k = m + k := by
-  induction k with grind
-
-@[grind =]
-theorem hyperoperation_two (m k : ℕ) : hyperoperation 2 m k = m * k := by
-  induction k with grind
-
-@[grind =]
-theorem hyperoperation_three (m k : ℕ) : hyperoperation 3 m k = m ^ k := by
-  induction k with grind [Nat.pow_succ] -- Ouch, this is a bad `grind` lemma.
-
-@[grind =]
-theorem hyperoperation_ge_three_one (n k : ℕ) : hyperoperation (n + 3) 1 k = 1 := by
-  induction n generalizing k with
-  | zero => grind
-  | succ n ih =>
-    cases k
-    · grind
-    · fail_if_success
-        -- This doesn't instantiate `hyperoperation_recursion`
-        -- because the goal contains `hyperoperation (n.succ + 3)`
-        -- but the lemma has `hyperoperation (n + 1)`.
-        grind [hyperoperation_recursion]
-      rw [hyperoperation_recursion, ih]
diff --git a/tests/lean/run/grind_bitvec2.lean b/tests/lean/run/grind_bitvec2.lean
index 5b887946ff..570cf4754d 100644
--- a/tests/lean/run/grind_bitvec2.lean
+++ b/tests/lean/run/grind_bitvec2.lean
@@ -457,7 +457,7 @@ theorem toInt_eq_toNat_of_lt {x : BitVec n} (h : 2 * x.toNat < 2^n) :
   grind [toInt_eq_toNat_cond]
 
 theorem msb_eq_false_iff_two_mul_lt {x : BitVec w} : x.msb = false ↔ 2 * x.toNat < 2^w := by
-  cases w <;> grind [Nat.pow_succ, msb_eq_decide]
+  cases w <;> grind [msb_eq_decide]
 
 grind_pattern msb_eq_false_iff_two_mul_lt => x.msb, x.toNat
 
diff --git a/tests/lean/run/grind_hyper_ex.lean b/tests/lean/run/grind_hyper_ex.lean
index 157f0b5a2c..9398f5a68c 100644
--- a/tests/lean/run/grind_hyper_ex.lean
+++ b/tests/lean/run/grind_hyper_ex.lean
@@ -28,8 +28,7 @@ theorem hyperoperation_two (m k : ℕ) : hyperoperation 2 m k = m * k := by
 
 @[grind =]
 theorem hyperoperation_three (m k : ℕ) : hyperoperation 3 m k = m ^ k := by
-  -- TODO: add support for Nat.pow_succ
-  induction k with grind [Nat.pow_succ] -- Ouch, this is a bad `grind` lemma.
+  induction k with grind
 
 @[grind =] theorem hyperoperation_ge_three_one (n k : ℕ) : hyperoperation (n + 3) 1 k = 1 := by
   induction n generalizing k with