1use super::FixedMontyParams;
2use super::mul::{almost_montgomery_reduce, mul_montgomery_form, square_repeat_montgomery_form};
3use crate::{AmmMultiplier, CtEq, Limb, MontyForm, Uint, UnsignedWithMontyForm, Word, word};
4use core::{array, mem};
5
6#[cfg(feature = "alloc")]
7use alloc::vec::Vec;
8
9const WINDOW: u32 = 4;
10const WINDOW_COUNT: usize = 1 << WINDOW;
11const WINDOW_MASK: Word = (1 << WINDOW) - 1;
12#[allow(clippy::integer_division_remainder_used, reason = "constant")]
13const BITS_PER_WINDOW: u32 = Limb::BITS / WINDOW;
14
15pub const fn pow_montgomery_form<
20 const LIMBS: usize,
21 const RHS_LIMBS: usize,
22 const VARTIME: bool,
23>(
24 x: &Uint<LIMBS>,
25 exponent: &Uint<RHS_LIMBS>,
26 exponent_bits: u32,
27 params: &FixedMontyParams<LIMBS>,
28) -> Uint<LIMBS> {
29 multi_exponentiate_montgomery_form_array::<LIMBS, RHS_LIMBS, 1, VARTIME>(
30 &[(*x, *exponent)],
31 exponent_bits,
32 params,
33 )
34}
35
36pub fn pow_montgomery_form_amm<'a, U>(
44 x: &U,
45 exponent: &U,
46 exponent_bits: u32,
47 params: &'a <U::MontyForm as MontyForm>::Params,
48) -> U
49where
50 U: UnsignedWithMontyForm,
51 <U::MontyForm as MontyForm>::Multiplier<'a>: AmmMultiplier<'a>,
52{
53 let exponent_bits = exponent_bits.min(exponent.bits_precision());
54
55 let one = params.as_ref().one().clone();
56
57 if exponent_bits == 0 {
58 return one; }
60
61 let mut multiplier = <U::MontyForm as MontyForm>::Multiplier::from(params);
62 let mut power = x.clone();
63
64 let powers: [U; WINDOW_COUNT] = array::from_fn(|n| {
66 if n == 0 {
67 one.clone()
68 } else if n == WINDOW_COUNT - 1 {
69 power.clone()
70 } else {
71 let mut new_power = power.clone();
72 multiplier.mul_amm_assign(&mut new_power, x);
73
74 mem::swap(&mut power, &mut new_power);
75 new_power
76 }
77 });
78
79 let (starting_limb, starting_window, starting_window_mask) = pow_init(exponent_bits);
80
81 let mut z = one; let mut power = powers[0].clone();
83
84 for limb_num in (0..=starting_limb).rev() {
85 let w = exponent.as_limbs()[limb_num].0;
86
87 let mut window_num = if limb_num == starting_limb {
88 starting_window + 1
89 } else {
90 BITS_PER_WINDOW
91 };
92
93 while window_num > 0 {
94 window_num -= 1;
95
96 let mut idx = (w >> (window_num * WINDOW)) & WINDOW_MASK;
97
98 if limb_num == starting_limb && window_num == starting_window {
99 idx &= starting_window_mask;
100 } else {
101 for _ in 1..=WINDOW {
102 multiplier.square_amm_assign(&mut z);
103 }
104 }
105
106 power.as_mut_limbs().copy_from_slice(powers[0].as_limbs());
108 for i in 1..WINDOW_COUNT {
109 power.ct_assign(&powers[i], (i as Word).ct_eq(&idx));
110 }
111
112 multiplier.mul_amm_assign(&mut z, &power);
113 }
114 }
115
116 almost_montgomery_reduce(z.as_mut(), params.as_ref().modulus().as_uint_ref());
117 z
118}
119
120pub const fn multi_exponentiate_montgomery_form_array<
121 const LIMBS: usize,
122 const RHS_LIMBS: usize,
123 const N: usize,
124 const VARTIME: bool,
125>(
126 bases_and_exponents: &[(Uint<LIMBS>, Uint<RHS_LIMBS>); N],
127 exponent_bits: u32,
128 params: &FixedMontyParams<LIMBS>,
129) -> Uint<LIMBS> {
130 let exponent_bits = if exponent_bits > Uint::<RHS_LIMBS>::BITS {
131 Uint::<RHS_LIMBS>::BITS
132 } else {
133 exponent_bits
134 };
135
136 if exponent_bits == 0 {
137 return *params.one(); }
139
140 let mut powers_and_exponents =
141 [([Uint::<LIMBS>::ZERO; WINDOW_COUNT], Uint::<RHS_LIMBS>::ZERO); N];
142
143 let mut i = 0;
144 while i < N {
145 let (base, exponent) = bases_and_exponents[i];
146 powers_and_exponents[i] = (compute_powers(&base, params), exponent);
147 i += 1;
148 }
149
150 multi_exponentiate_montgomery_form_internal::<LIMBS, RHS_LIMBS, VARTIME>(
151 &powers_and_exponents,
152 exponent_bits,
153 params,
154 )
155}
156
157#[cfg(feature = "alloc")]
164pub fn multi_exponentiate_montgomery_form_slice<
165 const LIMBS: usize,
166 const RHS_LIMBS: usize,
167 const VARTIME: bool,
168>(
169 bases_and_exponents: &[(Uint<LIMBS>, Uint<RHS_LIMBS>)],
170 exponent_bits: u32,
171 params: &FixedMontyParams<LIMBS>,
172) -> Uint<LIMBS> {
173 let exponent_bits = exponent_bits.min(Uint::<RHS_LIMBS>::BITS);
174
175 if exponent_bits == 0 {
176 return *params.one(); }
178
179 let powers_and_exponents: Vec<([Uint<LIMBS>; WINDOW_COUNT], Uint<RHS_LIMBS>)> =
180 bases_and_exponents
181 .iter()
182 .map(|(base, exponent)| (compute_powers(base, params), *exponent))
183 .collect();
184
185 multi_exponentiate_montgomery_form_internal::<LIMBS, RHS_LIMBS, VARTIME>(
186 powers_and_exponents.as_slice(),
187 exponent_bits,
188 params,
189 )
190}
191
192const fn compute_powers<const LIMBS: usize>(
193 x: &Uint<LIMBS>,
194 params: &FixedMontyParams<LIMBS>,
195) -> [Uint<LIMBS>; WINDOW_COUNT] {
196 let mut powers = [*params.one(); WINDOW_COUNT];
198 powers[1] = *x;
199
200 let mut i = 2;
201 while i < powers.len() {
202 powers[i] = mul_montgomery_form(&powers[i - 1], x, params.modulus(), params.mod_neg_inv());
203 i += 1;
204 }
205
206 powers
207}
208
209#[allow(clippy::cast_possible_truncation)]
210const fn multi_exponentiate_montgomery_form_internal<
211 const LIMBS: usize,
212 const RHS_LIMBS: usize,
213 const VARTIME: bool,
214>(
215 powers_and_exponents: &[([Uint<LIMBS>; WINDOW_COUNT], Uint<RHS_LIMBS>)],
216 exponent_bits: u32,
217 params: &FixedMontyParams<LIMBS>,
218) -> Uint<LIMBS> {
219 let (starting_limb, starting_window, starting_window_mask) = pow_init(exponent_bits);
220
221 let mut z = *params.one(); let mut limb_num = starting_limb + 1;
224 while limb_num > 0 {
225 limb_num -= 1;
226
227 let mut window_num = if limb_num == starting_limb {
228 starting_window + 1
229 } else {
230 BITS_PER_WINDOW
231 };
232 while window_num > 0 {
233 window_num -= 1;
234
235 if limb_num != starting_limb || window_num != starting_window {
236 z = square_repeat_montgomery_form(
237 &z,
238 WINDOW,
239 params.modulus(),
240 params.mod_neg_inv(),
241 );
242 }
243
244 let mut i = 0;
245 while i < powers_and_exponents.len() {
246 let (powers, exponent) = &powers_and_exponents[i];
247 let w = exponent.as_limbs()[limb_num].0;
248 let mut idx = (w >> (window_num * WINDOW)) & WINDOW_MASK;
249
250 if limb_num == starting_limb && window_num == starting_window {
251 idx &= starting_window_mask;
252 }
253
254 if VARTIME {
255 if idx > 0 {
256 z = mul_montgomery_form(
257 &z,
258 &powers[idx as usize],
259 params.modulus(),
260 params.mod_neg_inv(),
261 );
262 }
263 } else {
264 let mut power = powers[0];
266 let mut j = 1;
267 while j < WINDOW_COUNT {
268 let choice = word::choice_from_eq(j as Word, idx);
269 power = Uint::<LIMBS>::select(&power, &powers[j], choice);
270 j += 1;
271 }
272 z = mul_montgomery_form(&z, &power, params.modulus(), params.mod_neg_inv());
273 }
274
275 i += 1;
276 }
277 }
278 }
279
280 z
281}
282
283#[allow(clippy::integer_division_remainder_used, reason = "public parameter")]
287const fn pow_init(exponent_bits: u32) -> (usize, u32, Word) {
288 let starting_limb = ((exponent_bits - 1) / Limb::BITS) as usize;
289 let starting_bit_in_limb = (exponent_bits - 1) % Limb::BITS;
290 let starting_window = starting_bit_in_limb / WINDOW;
291 let starting_window_mask = (1 << (starting_bit_in_limb % WINDOW + 1)) - 1;
292 (starting_limb, starting_window, starting_window_mask)
293}
294
295#[cfg(test)]
296mod tests {
297 use super::{
298 multi_exponentiate_montgomery_form_array, pow_montgomery_form, pow_montgomery_form_amm,
299 };
300 use crate::{Odd, U128, modular::FixedMontyParams};
301
302 #[cfg(feature = "alloc")]
303 use super::multi_exponentiate_montgomery_form_slice;
304
305 const PARAMS: FixedMontyParams<{ U128::LIMBS }> =
306 FixedMontyParams::new_vartime(Odd::new_unchecked(U128::from_u8(101)));
307
308 #[test]
309 fn pow_montgomery_form_clamps_oversized_exponent_bits() {
310 let base = U128::from_u8(7);
311 let exponent = U128::from_u8(3);
312
313 let expected = pow_montgomery_form::<{ U128::LIMBS }, { U128::LIMBS }, false>(
314 &base,
315 &exponent,
316 U128::BITS,
317 &PARAMS,
318 );
319 let actual = pow_montgomery_form::<{ U128::LIMBS }, { U128::LIMBS }, false>(
320 &base,
321 &exponent,
322 U128::BITS + 1,
323 &PARAMS,
324 );
325
326 assert_eq!(actual, expected);
327 }
328
329 #[test]
330 fn pow_montgomery_form_amm_clamps_oversized_exponent_bits() {
331 let base = U128::from_u8(7);
332 let exponent = U128::from_u8(3);
333
334 let expected = pow_montgomery_form_amm(&base, &exponent, U128::BITS, &PARAMS);
335 let actual = pow_montgomery_form_amm(&base, &exponent, U128::BITS + 1, &PARAMS);
336
337 assert_eq!(actual, expected);
338 }
339
340 #[test]
341 fn multi_exponentiate_montgomery_form_array_clamps_oversized_exponent_bits() {
342 let bases_and_exponents = [(U128::from_u8(7), U128::from_u8(3))];
343
344 let expected = multi_exponentiate_montgomery_form_array::<
345 { U128::LIMBS },
346 { U128::LIMBS },
347 1,
348 false,
349 >(&bases_and_exponents, U128::BITS, &PARAMS);
350 let actual = multi_exponentiate_montgomery_form_array::<
351 { U128::LIMBS },
352 { U128::LIMBS },
353 1,
354 false,
355 >(&bases_and_exponents, U128::BITS + 1, &PARAMS);
356
357 assert_eq!(actual, expected);
358 }
359
360 #[cfg(feature = "alloc")]
361 #[test]
362 fn multi_exponentiate_montgomery_form_slice_clamps_oversized_exponent_bits() {
363 let bases_and_exponents = [(U128::from_u8(7), U128::from_u8(3))];
364
365 let expected = multi_exponentiate_montgomery_form_slice::<
366 { U128::LIMBS },
367 { U128::LIMBS },
368 false,
369 >(&bases_and_exponents, U128::BITS, &PARAMS);
370 let actual = multi_exponentiate_montgomery_form_slice::<
371 { U128::LIMBS },
372 { U128::LIMBS },
373 false,
374 >(&bases_and_exponents, U128::BITS + 1, &PARAMS);
375
376 assert_eq!(actual, expected);
377 }
378}