@@ -59,24 +59,25 @@ so we use an additional trick to do binary subdivision on `log2 b`. As a result
5959a proof of depth `log (log b)` which will essentially never overflow before the numbers involved
6060themselves exceed memory limits.
6161-/
62- partial def evalNatPow (a b : Q(ℕ)) : ( c : Q(ℕ)) × Q(Nat.pow $a $b = $c) :=
62+ partial def evalNatPow (a b : Q(ℕ)) : OptionT CoreM (( c : Q(ℕ)) × Q(Nat.pow $a $b = $c)) := do
6363 if b.natLit! = 0 then
6464 haveI : $b =Q 0 := ⟨⟩
65- ⟨q(nat_lit 1 ), q(natPow_zero)⟩
65+ return ⟨q(nat_lit 1 ), q(natPow_zero)⟩
6666 else if a.natLit! = 0 then
6767 haveI : $a =Q 0 := ⟨⟩
6868 have b' : Q(ℕ) := mkRawNatLit (b.natLit! - 1 )
6969 haveI : $b =Q Nat.succ $b' := ⟨⟩
70- ⟨q(nat_lit 0 ), q(zero_natPow)⟩
70+ return ⟨q(nat_lit 0 ), q(zero_natPow)⟩
7171 else if a.natLit! = 1 then
7272 haveI : $a =Q 1 := ⟨⟩
73- ⟨q(nat_lit 1 ), q(one_natPow)⟩
73+ return ⟨q(nat_lit 1 ), q(one_natPow)⟩
7474 else if b.natLit! = 1 then
7575 haveI : $b =Q 1 := ⟨⟩
76- ⟨a, q(natPow_one)⟩
76+ return ⟨a, q(natPow_one)⟩
7777 else
78+ guard <| ← Lean.checkExponent b.natLit!
7879 let ⟨c, p⟩ := go b.natLit!.log2 a q(nat_lit 1 ) a b _ .rfl
79- ⟨c, q(($p).run)⟩
80+ return ⟨c, q(($p).run)⟩
8081where
8182 /-- Invariants: `a ^ b₀ = c₀`, `depth > 0`, `b >>> depth = b₀`, `p := Nat.pow $a $b₀ = $c₀` -/
8283 go (depth : Nat) (a b₀ c₀ b : Q(ℕ)) (p : Q(Prop )) (hp : $p =Q (Nat.pow $a $b₀ = $c₀)) :
@@ -120,28 +121,29 @@ theorem intPow_negOfNat_bit1 {b' c' : ℕ} (h1 : Nat.pow a b' = c')
120121 simp [mul_comm, mul_left_comm]
121122
122123/-- Evaluates `Int.pow a b = c` where `a` and `b` are raw integer literals. -/
123- partial def evalIntPow (za : ℤ) (a : Q(ℤ)) (b : Q(ℕ)) : ℤ × (c : Q(ℤ)) × Q(Int.pow $a $b = $c) :=
124+ partial def evalIntPow (za : ℤ) (a : Q(ℤ)) (b : Q(ℕ)) :
125+ OptionT CoreM (ℤ × (c : Q(ℤ)) × Q(Int.pow $a $b = $c)) := do
124126 have a' : Q(ℕ) := a.appArg!
125127 if 0 ≤ za then
126- haveI : $a =Q .ofNat $a' := ⟨⟩
127- let ⟨c, p⟩ := evalNatPow a' b
128- ⟨c.natLit!, q(.ofNat $c), q(intPow_ofNat $p)⟩
128+ have : $a =Q .ofNat $a' := ⟨⟩
129+ let ⟨c, p⟩ ← evalNatPow a' b
130+ return ⟨c.natLit!, q(.ofNat $c), q(intPow_ofNat $p)⟩
129131 else
130- haveI : $a =Q .negOfNat $a' := ⟨⟩
132+ have : $a =Q .negOfNat $a' := ⟨⟩
131133 let b' := b.natLit!
132134 have b₀ : Q(ℕ) := mkRawNatLit (b' >>> 1 )
133- let ⟨c₀, p⟩ := evalNatPow a' b₀
135+ let ⟨c₀, p⟩ := ← evalNatPow a' b₀
134136 let c' := c₀.natLit!
135137 if b' &&& 1 == 0 then
136138 have c : Q(ℕ) := mkRawNatLit (c' * c')
137139 have pc : Q($c₀ * $c₀ = $c) := (q(Eq.refl $c) : Expr)
138140 have pb : Q(2 * $b₀ = $b) := (q(Eq.refl $b) : Expr)
139- ⟨c.natLit!, q(.ofNat $c), q(intPow_negOfNat_bit0 $p $pb $pc)⟩
141+ return ⟨c.natLit!, q(.ofNat $c), q(intPow_negOfNat_bit0 $p $pb $pc)⟩
140142 else
141143 have c : Q(ℕ) := mkRawNatLit (c' * (c' * a'.natLit!))
142144 have pc : Q($c₀ * ($c₀ * $a') = $c) := (q(Eq.refl $c) : Expr)
143145 have pb : Q(2 * $b₀ + 1 = $b) := (q(Eq.refl $b) : Expr)
144- ⟨-c.natLit!, q(.negOfNat $c), q(intPow_negOfNat_bit1 $p $pb $pc)⟩
146+ return ⟨-c.natLit!, q(.negOfNat $c), q(intPow_negOfNat_bit1 $p $pb $pc)⟩
145147
146148-- see note [norm_num lemma function equality]
147149theorem isNat_pow {α} [Semiring α] : ∀ {f : α → ℕ → α} {a : α} {b a' b' c : ℕ},
@@ -175,30 +177,30 @@ theorem isNNRat_pow {α} [Semiring α] {f : α → ℕ → α} {a : α} {an cn :
175177/-- Main part of `evalPow`. -/
176178def evalPow.core {u : Level} {α : Q(Type u)} (e : Q(«$α »)) (f : Q(«$α » → ℕ → «$α »)) (a : Q(«$α »))
177179 (b nb : Q(ℕ)) (pb : Q(IsNat «$b » «$nb »)) (sα : Q(Semiring «$α »)) (ra : Result a) :
178- Option (Result e) := do
180+ OptionT CoreM (Result e) := do
179181 haveI' : $e =Q $a ^ $b := ⟨⟩
180182 haveI' : $f =Q HPow.hPow := ⟨⟩
181183 match ra with
182184 | .isBool .. => failure
183185 | .isNat sα na pa =>
184186 assumeInstancesCommute
185- have ⟨c, r⟩ := evalNatPow na nb
187+ let ⟨c, r⟩ ← evalNatPow na nb
186188 return .isNat sα c q(isNat_pow (f := $f) (.refl $f) $pa $pb $r)
187189 | .isNegNat rα .. =>
188190 assumeInstancesCommute
189- let ⟨za, na, pa⟩ ← ra.toInt rα
190- have ⟨zc, c, r⟩ := evalIntPow za na nb
191+ let ⟨za, na, pa⟩ ← OptionT.mk <| pure ( ra.toInt rα)
192+ let ⟨zc, c, r⟩ ← evalIntPow za na nb
191193 return .isInt rα c zc q(isInt_pow (f := $f) (.refl $f) $pa $pb $r)
192194 | .isNNRat dα _qa na da pa =>
193195 assumeInstancesCommute
194- have ⟨nc, r1⟩ := evalNatPow na nb
195- have ⟨dc, r2⟩ := evalNatPow da nb
196+ let ⟨nc, r1⟩ ← evalNatPow na nb
197+ let ⟨dc, r2⟩ ← evalNatPow da nb
196198 let qc := mkRat nc.natLit! dc.natLit!
197199 return .isNNRat dα qc nc dc q(isNNRat_pow (f := $f) (.refl $f) $pa $pb $r1 $r2)
198200 | .isNegNNRat dα qa na da pa =>
199201 assumeInstancesCommute
200- have ⟨zc, nc, r1⟩ := evalIntPow qa.num q(Int.negOfNat $na) nb
201- have ⟨dc, r2⟩ := evalNatPow da nb
202+ let ⟨zc, nc, r1⟩ ← evalIntPow qa.num q(Int.negOfNat $na) nb
203+ let ⟨dc, r2⟩ ← evalNatPow da nb
202204 let qc := mkRat zc dc.natLit!
203205 return .isRat dα qc nc dc q(isRat_pow (f := $f) (.refl $f) $pa $pb $r1 $r2)
204206
@@ -214,7 +216,9 @@ def evalPow : NormNumExt where eval {u α} e := do
214216 guard <|← withDefault <| withNewMCtxDepth <| isDefEq f q(HPow.hPow (α := $α))
215217 haveI' : $e =Q $a ^ $b := ⟨⟩
216218 haveI' : $f =Q HPow.hPow := ⟨⟩
217- evalPow.core q($e) q($f) q($a) q($b) q($nb) q($pb) q($sα) ra
219+ let .some r ←
220+ liftM <| OptionT.run (evalPow.core q($e) q($f) q($a) q($b) q($nb) q($pb) q($sα) ra) | failure
221+ return r
218222
219223theorem isNat_zpow_pos {α : Type *} [DivisionSemiring α] {a : α} {b : ℤ} {nb ne : ℕ}
220224 (pb : IsNat b nb) (pe' : IsNat (a ^ nb) ne) :
0 commit comments