1use super::Uint;
4use crate::{Checked, Choice, CtOption, Limb, PowBoundedExp, Wrapping};
5
6impl<const LIMBS: usize> Uint<LIMBS> {
7 #[must_use]
9 pub const fn checked_pow<const RHS_LIMBS: usize>(
10 &self,
11 exp: &Uint<RHS_LIMBS>,
12 ) -> CtOption<Self> {
13 self.checked_pow_bounded_exp(exp, Uint::<RHS_LIMBS>::BITS)
14 }
15
16 #[must_use]
23 pub const fn checked_pow_bounded_exp<const RHS_LIMBS: usize>(
24 &self,
25 exp: &Uint<RHS_LIMBS>,
26 exp_bits: u32,
27 ) -> CtOption<Self> {
28 let (lo, overflow) = self.overflowing_pow_bounded_exp(exp, exp_bits);
29 CtOption::new(lo, overflow.not())
30 }
31
32 #[must_use]
36 pub const fn checked_pow_vartime<const RHS_LIMBS: usize>(
37 &self,
38 exp: &Uint<RHS_LIMBS>,
39 ) -> CtOption<Self> {
40 let (lo, overflow) = self.overflowing_pow_vartime(exp);
41 CtOption::new(lo, overflow.not())
42 }
43
44 #[must_use]
46 pub const fn saturating_pow<const RHS_LIMBS: usize>(&self, exp: &Uint<RHS_LIMBS>) -> Self {
47 self.saturating_pow_bounded_exp(exp, Uint::<RHS_LIMBS>::BITS)
48 }
49
50 #[must_use]
57 pub const fn saturating_pow_bounded_exp<const RHS_LIMBS: usize>(
58 &self,
59 exp: &Uint<RHS_LIMBS>,
60 exp_bits: u32,
61 ) -> Self {
62 let (lo, overflow) = self.overflowing_pow_bounded_exp(exp, exp_bits);
63 Self::select(&lo, &Self::MAX, overflow)
64 }
65
66 #[must_use]
70 pub const fn saturating_pow_vartime<const RHS_LIMBS: usize>(
71 &self,
72 exp: &Uint<RHS_LIMBS>,
73 ) -> Self {
74 let (lo, overflow) = self.overflowing_pow_vartime(exp);
75 Self::select(&lo, &Self::MAX, overflow)
76 }
77
78 #[must_use]
80 pub const fn wrapping_pow<const RHS_LIMBS: usize>(&self, exp: &Uint<RHS_LIMBS>) -> Self {
81 self.wrapping_pow_bounded_exp(exp, Uint::<RHS_LIMBS>::BITS)
82 }
83
84 #[must_use]
91 pub const fn wrapping_pow_bounded_exp<const RHS_LIMBS: usize>(
92 &self,
93 exp: &Uint<RHS_LIMBS>,
94 exp_bits: u32,
95 ) -> Self {
96 assert!(exp_bits <= Uint::<RHS_LIMBS>::BITS);
97 let mut ret = Self::ONE;
98 let mut limbs = exp_bits.div_ceil(Limb::BITS);
99 let mut i = exp_bits - limbs.saturating_sub(1) * Limb::BITS;
100 while limbs > 0 {
101 limbs -= 1;
102 let limb = exp.limbs[limbs as usize];
103 while i > 0 {
104 i -= 1;
105 ret = ret.wrapping_square();
106 let apply = limb.shr(i).lsb_to_choice();
107 ret = Uint::select(&ret, &ret.wrapping_mul(self), apply);
108 }
109 i = Limb::BITS;
110 }
111 ret
112 }
113
114 #[inline]
118 #[must_use]
119 pub const fn wrapping_pow_vartime<const RHS_LIMBS: usize>(
120 &self,
121 exp: &Uint<RHS_LIMBS>,
122 ) -> Self {
123 let mut exp_bits = exp.bits_vartime();
124 if exp_bits == 0 {
125 return Self::ONE;
126 }
127 exp_bits -= 1;
128 let mut ret = *self;
129 let mut limbs = exp_bits.div_ceil(Limb::BITS);
130 let mut i = exp_bits - limbs.saturating_sub(1) * Limb::BITS;
131 while limbs > 0 {
132 limbs -= 1;
133 let word = exp.limbs[limbs as usize].0;
134 while i > 0 {
135 i -= 1;
136 ret = ret.wrapping_square();
137 if (word >> i) & 1 == 1 {
138 ret = ret.wrapping_mul(self);
139 }
140 }
141 i = Limb::BITS;
142 }
143 ret
144 }
145
146 #[inline(always)]
151 #[must_use]
152 pub(crate) const fn overflowing_pow_bounded_exp<const RHS_LIMBS: usize>(
153 &self,
154 exp: &Uint<RHS_LIMBS>,
155 exp_bits: u32,
156 ) -> (Self, Choice) {
157 assert!(exp_bits <= Uint::<RHS_LIMBS>::BITS);
158 let mut ret = Self::ONE;
159 let mut overflow = Choice::FALSE;
160 let mut check;
161 let mut limbs = exp_bits.div_ceil(Limb::BITS);
162 let mut i = exp_bits - limbs.saturating_sub(1) * Limb::BITS;
163 while limbs > 0 {
164 limbs -= 1;
165 let limb = exp.limbs[limbs as usize];
166 while i > 0 {
167 i -= 1;
168 (ret, check) = ret.overflowing_square();
169 overflow = overflow.or(check);
170 let (mul, check) = ret.overflowing_mul(self);
171 let apply = limb.shr(i).lsb_to_choice();
172 ret = Uint::select(&ret, &mul, apply);
173 overflow = overflow.or(check.and(apply));
174 }
175 i = Limb::BITS;
176 }
177 (ret, overflow)
178 }
179
180 #[inline(always)]
185 #[must_use]
186 pub(crate) const fn overflowing_pow_vartime<const RHS_LIMBS: usize>(
187 &self,
188 exp: &Uint<RHS_LIMBS>,
189 ) -> (Self, Choice) {
190 let mut exp_bits = exp.bits_vartime();
191 if exp_bits == 0 {
192 return (Self::ONE, Choice::FALSE);
193 }
194 exp_bits -= 1;
195 let mut ret = *self;
196 let mut overflow = Choice::FALSE;
197 let mut check;
198 let mut limbs = exp_bits.div_ceil(Limb::BITS);
199 let mut i = exp_bits - limbs.saturating_sub(1) * Limb::BITS;
200 while limbs > 0 {
201 limbs -= 1;
202 let word = exp.limbs[limbs as usize].0;
203 while i > 0 {
204 i -= 1;
205 (ret, check) = ret.overflowing_square();
206 overflow = overflow.or(check);
207 if (word >> i) & 1 == 1 {
208 (ret, check) = ret.overflowing_mul(self);
209 overflow = overflow.or(check);
210 }
211 }
212 i = Limb::BITS;
213 }
214 (ret, overflow)
215 }
216}
217
218impl<const LIMBS: usize, const RHS_LIMBS: usize> PowBoundedExp<Uint<RHS_LIMBS>>
219 for Checked<Uint<LIMBS>>
220{
221 fn pow_bounded_exp(&self, exponent: &Uint<RHS_LIMBS>, exponent_bits: u32) -> Self {
222 Checked(
223 self.0
224 .and_then(|base| base.checked_pow_bounded_exp(exponent, exponent_bits)),
225 )
226 }
227}
228
229impl<const LIMBS: usize, const RHS_LIMBS: usize> PowBoundedExp<Uint<RHS_LIMBS>>
230 for Wrapping<Uint<LIMBS>>
231{
232 fn pow_bounded_exp(&self, exponent: &Uint<RHS_LIMBS>, exponent_bits: u32) -> Self {
233 Wrapping(self.0.wrapping_pow_bounded_exp(exponent, exponent_bits))
234 }
235}
236
237#[cfg(test)]
238mod tests {
239 use crate::{Checked, CtOption, Pow, U128, Wrapping};
240
241 #[test]
242 fn checked_pow_expected() {
243 let checks = [
244 (U128::ZERO, U128::ZERO, Some(U128::ONE)),
245 (U128::ZERO, U128::MAX, Some(U128::ZERO)),
246 (U128::ONE, U128::ZERO, Some(U128::ONE)),
247 (U128::ONE, U128::MAX, Some(U128::ONE)),
248 (U128::MAX, U128::ZERO, Some(U128::ONE)),
249 (U128::MAX, U128::ONE, Some(U128::MAX)),
250 (U128::MAX, U128::from_u8(2), None),
251 (U128::MAX, U128::MAX, None),
252 ];
253 for (base, pow, expect) in checks {
254 assert_eq!(base.checked_pow(&pow).into_option(), expect);
255 assert_eq!(base.checked_pow_vartime(&pow).into_option(), expect);
256 assert_eq!(
257 Checked(CtOption::some(base)).pow(&pow).0.into_option(),
258 expect
259 );
260 }
261
262 assert!(
263 Checked(CtOption::<U128>::none())
264 .pow(&U128::ONE)
265 .0
266 .is_none()
267 .to_bool_vartime()
268 );
269 }
270
271 #[test]
272 fn saturating_pow_expected() {
273 let checks = [
274 (U128::ZERO, U128::ZERO, U128::ONE),
275 (U128::ZERO, U128::MAX, U128::ZERO),
276 (U128::ONE, U128::ZERO, U128::ONE),
277 (U128::ONE, U128::MAX, U128::ONE),
278 (U128::MAX, U128::ZERO, U128::ONE),
279 (U128::MAX, U128::ONE, U128::MAX),
280 (U128::MAX, U128::from_u8(2), U128::MAX),
281 (U128::MAX, U128::MAX, U128::MAX),
282 ];
283 for (base, pow, expect) in checks {
284 assert_eq!(base.saturating_pow(&pow), expect);
285 assert_eq!(base.saturating_pow_vartime(&pow), expect);
286 }
287 }
288
289 #[test]
290 fn wrapping_pow_expected() {
291 let checks = [
292 (U128::ZERO, U128::ZERO, U128::ONE),
293 (U128::ZERO, U128::MAX, U128::ZERO),
294 (U128::ONE, U128::ZERO, U128::ONE),
295 (U128::ONE, U128::MAX, U128::ONE),
296 (U128::MAX, U128::ZERO, U128::ONE),
297 (U128::MAX, U128::ONE, U128::MAX),
298 (U128::MAX, U128::from_u8(2), U128::ONE),
299 (U128::MAX, U128::from_u8(3), U128::MAX),
300 (U128::MAX, U128::MAX, U128::MAX),
301 ];
302 for (base, pow, expect) in checks {
303 assert_eq!(base.wrapping_pow(&pow), expect);
304 assert_eq!(base.wrapping_pow_vartime(&pow), expect);
305 assert_eq!(Wrapping(base).pow(&pow), Wrapping(expect));
306 }
307
308 let two = U128::from_u8(2);
309 assert_eq!(two.wrapping_pow_vartime(&U128::ZERO), U128::ONE);
310 for i in 0..10 {
311 let pow = 1u32 << i;
312 assert_eq!(
313 two.wrapping_pow_vartime(&U128::from_u32(pow)),
314 U128::from_u128(2u128.wrapping_pow(pow)),
315 "i={i}"
316 );
317 }
318
319 assert_eq!(
320 U128::from_u8(3).wrapping_pow_vartime(&U128::from_u128((1 << 64) + (1 << 63))),
321 U128::from_be_hex("002b3854b3dc5d6e0000000000000001")
322 );
323 }
324}