ark_ff_macros/montgomery/
mod.rs

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