feat: implement zpow norm_num extension (#9875)

Co-authored-by: Eric Wieser <[email protected]>

Mathlib/Tactic/NormNum/Pow.lean

@@ -193,3 +193,73 @@ def evalPow : NormNumExt where eval {u α} e := do
       let qc := mkRat zc dc.natLit!
       return .isRat' dα qc nc dc q(isRat_pow (f := $f) (.refl $f) $pa $pb $r1 $r2)
   core
+
+theorem isNat_zpow_pos {α : Type*} [DivisionSemiring α] {a : α} {b : ℤ} {nb ne : ℕ}
+    (pb : IsNat b nb) (pe' : IsNat (a ^ nb) ne) :
+    IsNat (a ^ b) ne := by
+  rwa [pb.out, zpow_coe_nat]
+
+theorem isNat_zpow_neg {α : Type*} [DivisionSemiring α] {a : α} {b : ℤ} {nb ne : ℕ}
+    (pb : IsInt b (Int.negOfNat nb)) (pe' : IsNat (a ^ nb)⁻¹ ne) :
+    IsNat (a ^ b) ne := by
+  rwa [pb.out, Int.cast_negOfNat, zpow_neg, zpow_coe_nat]
+
+theorem isInt_zpow_pos {α : Type*} [DivisionRing α] {a : α} {b : ℤ} {nb ne : ℕ}
+    (pb : IsNat b nb) (pe' : IsInt (a ^ nb) (Int.negOfNat ne)) :
+    IsInt (a ^ b) (Int.negOfNat ne) := by
+  rwa [pb.out, zpow_coe_nat]
+
+theorem isInt_zpow_neg {α : Type*} [DivisionRing α] {a : α} {b : ℤ} {nb ne : ℕ}
+    (pb : IsInt b (Int.negOfNat nb)) (pe' : IsInt (a ^ nb)⁻¹ (Int.negOfNat ne)) :
+    IsInt (a ^ b) (Int.negOfNat ne) := by
+  rwa [pb.out, Int.cast_negOfNat, zpow_neg, zpow_coe_nat]
+
+theorem isRat_zpow_pos {α : Type*} [DivisionRing α] {a : α} {b : ℤ} {nb : ℕ}
+    {num : ℤ} {den : ℕ}
+    (pb : IsNat b nb) (pe' : IsRat (a^nb) num den) :
+    IsRat (a^b) num den := by
+  rwa [pb.out, zpow_coe_nat]
+
+theorem isRat_zpow_neg {α : Type*} [DivisionRing α] {a : α} {b : ℤ} {nb : ℕ}
+    {num : ℤ} {den : ℕ}
+    (pb : IsInt b (Int.negOfNat nb)) (pe' : IsRat ((a^nb)⁻¹) num den) :
+    IsRat (a^b) num den := by
+  rwa [pb.out, Int.cast_negOfNat, zpow_neg, zpow_coe_nat]
+
+/-- The `norm_num` extension which identifies expressions of the form `a ^ b`,
+such that `norm_num` successfully recognises both `a` and `b`, with `b : ℤ`. -/
+@[norm_num (_ : α) ^ (_ : ℤ)]
+def evalZPow : NormNumExt where eval {u α} e := do
+  let .app (.app (f : Q($α → ℤ → $α)) (a : Q($α))) (b : Q(ℤ)) ← whnfR e | failure
+  let _c ← synthInstanceQ q(DivisionSemiring $α)
+  let rb ← derive (α := q(ℤ)) b
+  have h : $e =Q $a ^ $b := ⟨⟩
+  h.check
+  match rb with
+  | .isBool .. | .isRat _ .. => failure
+  | .isNat sβ nb pb =>
+    match ← derive q($a ^ $nb) with
+    | .isBool .. => failure
+    | .isNat sα' ne' pe' =>
+      assumeInstancesCommute
+      return .isNat sα' ne' q(isNat_zpow_pos $pb $pe')
+    | .isNegNat sα' ne' pe' =>
+      let _c ← synthInstanceQ q(DivisionRing $α)
+      assumeInstancesCommute
+      return .isNegNat sα' ne' q(isInt_zpow_pos $pb $pe')
+    | .isRat sα' qe' nume' dene' pe' =>
+      assumeInstancesCommute
+      return .isRat sα' qe' nume' dene' q(isRat_zpow_pos $pb $pe')
+  | .isNegNat sβ nb pb =>
+    match ← derive q(($a ^ $nb)⁻¹) with
+    | .isBool .. => failure
+    | .isNat sα' ne' pe' =>
+      assumeInstancesCommute
+      return .isNat sα' ne' q(isNat_zpow_neg $pb $pe')
+    | .isNegNat sα' ne' pe' =>
+      let _c ← synthInstanceQ q(DivisionRing $α)
+      assumeInstancesCommute
+      return .isNegNat sα' ne' q(isInt_zpow_neg $pb $pe')
+    | .isRat sα' qe' nume' dene' pe' =>
+      assumeInstancesCommute
+      return .isRat sα' qe' nume' dene' q(isRat_zpow_neg $pb $pe')

test/norm_num.lean

@@ -17,6 +17,10 @@ axiom Real : Type
 notation "ℝ" => Real
 @[instance] axiom Real.linearOrderedRing : LinearOrderedField ℝ
 
+axiom NNReal : Type
+notation "ℝ≥0" => NNReal
+@[instance] axiom NNReal.linearOrderedsemifield : LinearOrderedSemifield ℝ≥0
+
 axiom Complex : Type
 notation "ℂ" => Complex
 @[instance] axiom Complex.field : Field ℂ
@@ -38,11 +42,15 @@ example : (7:ℝ)/2 > 3 := by norm_num1
 example : (4:ℝ)⁻¹ < 1 := by norm_num1
 example : ((1:ℝ) / 2)⁻¹ = 2 := by norm_num1
 example : 2 ^ 17 - 1 = 131071 := by norm_num1
--- example : (3 : ℝ) ^ (-2 : ℤ) = 1/9 := by norm_num1
--- example : (3 : ℝ) ^ (-2 : ℤ) = 1/9 := by norm_num1
--- example : (-3 : ℝ) ^ (0 : ℤ) = 1 := by norm_num1
--- example : (-3 : ℝ) ^ (-1 : ℤ) = -1/3 := by norm_num1
--- example : (-3 : ℝ) ^ (2 : ℤ) = 9 := by norm_num1
+example : (3 : ℝ) ^ (-2 : ℤ) = 1/9 := by norm_num1
+example : (-3 : ℝ) ^ (0 : ℤ) = 1 := by norm_num1
+example : (-3 : ℝ) ^ (-1 : ℤ) = -1/3 := by norm_num1
+example : (-3 : ℝ) ^ (1 : ℤ) = -3 := by norm_num1
+example : (-3 : ℝ) ^ (2 : ℤ) = 9 := by norm_num1
+example : (1/3 : ℝ) ^ (2 : ℤ) = 1/9 := by norm_num1
+example : (1/3 : ℝ) ^ (-2 : ℤ) = 9 := by norm_num1
+example : (-1/3 : ℝ) ^ (-1 : ℤ) = -3 := by norm_num1
+example : (3 : ℝ≥0) ^ (2 : ℤ) = 9 := by norm_num1
 
 section InvLit