ark_ff_macros/montgomery/
mod.rs

1use std::str::FromStr;
2
3use num_bigint::BigUint;
4use num_traits::One;
5
6mod biginteger;
7use biginteger::{add_with_carry_impl, sub_with_borrow_impl, subtract_modulus_impl};
8
9mod add;
10use add::add_assign_impl;
11mod double;
12use double::double_in_place_impl;
13mod mul;
14use mul::mul_assign_impl;
15
16mod square;
17use square::square_in_place_impl;
18
19mod sum_of_products;
20use sum_of_products::sum_of_products_impl;
21
22use crate::utils;
23
24pub fn mont_config_helper(
25    modulus: BigUint,
26    generator: BigUint,
27    small_subgroup_base: Option<u32>,
28    small_subgroup_power: Option<u32>,
29    config_name: proc_macro2::Ident,
30) -> proc_macro2::TokenStream {
31    let mut limbs = 1usize;
32    {
33        let mut cur = BigUint::one() << 64; // always 64-bit limbs for now
34        while cur < modulus {
35            limbs += 1;
36            cur <<= 64;
37        }
38    }
39
40    // modulus - 1 = 2^s * t
41    let mut trace = &modulus - BigUint::from_str("1").unwrap();
42    while !trace.bit(0) {
43        trace >>= 1u8;
44    }
45
46    // Compute 2^s root of unity given the generator
47    let remaining_subgroup_size = match (small_subgroup_base, small_subgroup_power) {
48        (Some(base), Some(power)) => Some(&trace / BigUint::from(base).pow(power)),
49        (None, None) => None,
50        (..) => panic!("Must specify both `small_subgroup_base` and `small_subgroup_power`"),
51    };
52    let two_adic_root_of_unity = generator.modpow(&trace, &modulus);
53    let large_subgroup_generator = remaining_subgroup_size
54        .as_ref()
55        .map(|e| generator.modpow(e, &modulus).to_string());
56    let modulus = modulus.to_string();
57    let generator = generator.to_string();
58    let two_adic_root_of_unity = two_adic_root_of_unity.to_string();
59    let modulus_limbs = utils::str_to_limbs_u64(&modulus).1;
60    let modulus_has_spare_bit = modulus_limbs.last().unwrap() >> 63 == 0;
61    let can_use_no_carry_mul_opt = {
62        let first_limb_check = *modulus_limbs.last().unwrap() < (u64::MAX >> 1);
63        if limbs != 1 {
64            first_limb_check && modulus_limbs[..limbs - 1].iter().any(|l| *l != u64::MAX)
65        } else {
66            first_limb_check
67        }
68    };
69    let modulus = quote::quote! { BigInt([ #( #modulus_limbs ),* ]) };
70
71    let add_with_carry = add_with_carry_impl(limbs);
72    let sub_with_borrow = sub_with_borrow_impl(limbs);
73    let subtract_modulus = subtract_modulus_impl(&modulus);
74    let add_assign = add_assign_impl(modulus_has_spare_bit);
75    let double_in_place = double_in_place_impl(modulus_has_spare_bit);
76    let mul_assign = mul_assign_impl(
77        can_use_no_carry_mul_opt,
78        limbs,
79        &modulus_limbs,
80        modulus_has_spare_bit,
81    );
82    let square_in_place = square_in_place_impl(
83        can_use_no_carry_mul_opt,
84        limbs,
85        &modulus_limbs,
86        modulus_has_spare_bit,
87    );
88    let sum_of_products = sum_of_products_impl(limbs, &modulus_limbs);
89
90    let mixed_radix = if let Some(large_subgroup_generator) = large_subgroup_generator {
91        quote::quote! {
92            const SMALL_SUBGROUP_BASE: Option<u32> = Some(#small_subgroup_base);
93
94            const SMALL_SUBGROUP_BASE_ADICITY: Option<u32> = Some(#small_subgroup_power);
95
96            const LARGE_SUBGROUP_ROOT_OF_UNITY: Option<F> = Some(ark_ff::MontFp!(#large_subgroup_generator));
97        }
98    } else {
99        quote::quote! {}
100    };
101
102    quote::quote! {
103        const _: () = {
104            use ark_ff::{fields::Fp, BigInt, BigInteger, biginteger::arithmetic as fa, fields::*};
105            type B = BigInt<#limbs>;
106            type F = Fp<MontBackend<#config_name, #limbs>, #limbs>;
107
108            #[automatically_derived]
109            impl MontConfig<#limbs> for #config_name {
110                const MODULUS: B = #modulus;
111
112                const GENERATOR: F = ark_ff::MontFp!(#generator);
113
114                const TWO_ADIC_ROOT_OF_UNITY: F = ark_ff::MontFp!(#two_adic_root_of_unity);
115
116                #mixed_radix
117
118                #[inline(always)]
119                fn add_assign(a: &mut F, b: &F) {
120                    #add_assign
121                }
122
123                #[inline(always)]
124                fn sub_assign(a: &mut F, b: &F) {
125                    // If `other` is larger than `self`, add the modulus to self first.
126                    if b.0 > a.0 {
127                        __add_with_carry(&mut a.0, &#modulus);
128                    }
129                    __sub_with_borrow(&mut a.0, &b.0);
130                }
131
132                #[inline(always)]
133                fn double_in_place(a: &mut F) {
134                    #double_in_place
135                }
136
137                /// Sets `a = -a`.
138                #[inline(always)]
139                fn neg_in_place(a: &mut F) {
140                    if *a != F::ZERO {
141                        let mut tmp = #modulus;
142                        __sub_with_borrow(&mut tmp, &a.0);
143                        a.0 = tmp;
144                    }
145                }
146
147                #[inline(always)]
148                fn mul_assign(a: &mut F, b: &F) {
149                    #mul_assign
150                }
151                #[inline(always)]
152                fn square_in_place(a: &mut F) {
153                    #square_in_place
154                }
155
156                fn sum_of_products<const M: usize>(
157                    a: &[F; M],
158                    b: &[F; M],
159                ) -> F {
160                    #sum_of_products
161                }
162            }
163
164            #subtract_modulus
165
166            #add_with_carry
167
168            #sub_with_borrow
169        };
170    }
171}