Ark-ff Small Field Support

CHANGELOG.md

@@ -11,6 +11,7 @@
 - (`ark-poly`) Add fast polynomial division
 - (`ark-ec`) Improve GLV scalar multiplication performance by skipping leading zeroes.
 - (`ark-poly`) Make `SparsePolynomial.coeffs` field public
+- [\#1044](https://github.com/arkworks-rs/algebra/pull/1044) Add implementation for small field with native integer types
 
 ### Breaking changes
 

ff-macros/src/lib.rs

@@ -12,6 +12,7 @@ use proc_macro::TokenStream;
 use syn::{Expr, ExprLit, Item, ItemFn, Lit, Meta};
 
 mod montgomery;
+mod small_fp;
 mod unroll;
 
 pub(crate) mod utils;
@@ -74,6 +75,34 @@ pub fn mont_config(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     .into()
 }
 
+/// Derive the `SmallFpConfig` trait for small prime fields.
+///
+/// The attributes available to this macro are:
+/// * `modulus`: Specify the prime modulus underlying this prime field.
+/// * `generator`: Specify the generator of the multiplicative subgroup.
+/// * `backend`: Specify either "standard" or "montgomery" backend.
+#[proc_macro_derive(SmallFpConfig, attributes(modulus, generator, backend))]
+pub fn small_fp_config(input: TokenStream) -> TokenStream {
+    let ast: syn::DeriveInput = syn::parse(input).unwrap();
+
+    let modulus: u128 = fetch_attr("modulus", &ast.attrs)
+        .expect("Please supply a modulus attribute")
+        .parse()
+        .expect("Modulus should be a number");
+
+    let generator: u128 = fetch_attr("generator", &ast.attrs)
+        .expect("Please supply a generator attribute")
+        .parse()
+        .expect("Generator should be a number");
+
+    let backend: String = fetch_attr("backend", &ast.attrs)
+        .expect("Please supply a backend attribute")
+        .parse()
+        .expect("Backend should be a string");
+
+    small_fp::small_fp_config_helper(modulus, generator, backend, ast.ident).into()
+}
+
 const ARG_MSG: &str = "Failed to parse unroll threshold; must be a positive integer";
 
 /// Attribute used to unroll for loops found inside a function block.

ff-macros/src/small_fp/mod.rs

@@ -0,0 +1,44 @@
+mod montgomery_backend;
+mod standard_backend;
+mod utils;
+
+use quote::quote;
+
+/// This function is called by the `#[derive(SmallFp)]` macro and generates
+/// the implementation of the `SmallFpConfig`
+pub(crate) fn small_fp_config_helper(
+    modulus: u128,
+    generator: u128,
+    backend: String,
+    config_name: proc_macro2::Ident,
+) -> proc_macro2::TokenStream {
+    let ty = match modulus {
+        m if m < 1u128 << 8 => quote! { u8 },
+        m if m < 1u128 << 16 => quote! { u16 },
+        m if m < 1u128 << 32 => quote! { u32 },
+        m if m < 1u128 << 64 => quote! { u64 },
+        _ => quote! { u128 },
+    };
+
+    let backend_impl = match backend.as_str() {
+        "standard" => standard_backend::backend_impl(&ty, modulus, generator),
+        "montgomery" => montgomery_backend::backend_impl(&ty, modulus, generator),
+        _ => panic!("Unknown backend type: {}", backend),
+    };
+
+    let new_impl = match backend.as_str() {
+        "standard" => standard_backend::new(),
+        "montgomery" => montgomery_backend::new(modulus, ty),
+        _ => panic!("Unknown backend type: {}", backend),
+    };
+
+    quote! {
+        impl SmallFpConfig for #config_name {
+            #backend_impl
+        }
+
+        impl #config_name {
+            #new_impl
+        }
+    }
+}

ff-macros/src/small_fp/montgomery_backend.rs

@@ -0,0 +1,258 @@
+use std::u32;
+
+use super::*;
+use crate::small_fp::utils::{
+    compute_two_adic_root_of_unity, compute_two_adicity, generate_montgomery_bigint_casts,
+    generate_sqrt_precomputation, mod_mul_const,
+};
+
+pub(crate) fn backend_impl(
+    ty: &proc_macro2::TokenStream,
+    modulus: u128,
+    generator: u128,
+) -> proc_macro2::TokenStream {
+    let k_bits = 128 - modulus.leading_zeros();
+    let r: u128 = 1u128 << k_bits;
+    let r_mod_n = r % modulus;
+    let r_mask = r - 1;
+
+    let n_prime = mod_inverse_pow2(modulus, k_bits);
+    let one_mont = r_mod_n;
+    let generator_mont = mod_mul_const(generator % modulus, r_mod_n % modulus, modulus);
+
+    let two_adicity = compute_two_adicity(modulus);
+    let two_adic_root = compute_two_adic_root_of_unity(modulus, two_adicity);
+    let two_adic_root_mont = mod_mul_const(two_adic_root, r_mod_n, modulus);
+
+    let neg_one_mont = mod_mul_const(modulus - 1, r_mod_n, modulus);
+
+    let (from_bigint_impl, into_bigint_impl) =
+        generate_montgomery_bigint_casts(modulus, k_bits, r_mod_n);
+    let sqrt_precomp_impl = generate_sqrt_precomputation(modulus, two_adicity);
+
+    // Generate multiplication implementation based on type
+    let mul_impl = generate_mul_impl(ty, modulus, k_bits, r_mask, n_prime);
+
+    quote! {
+        type T = #ty;
+        const MODULUS: Self::T = #modulus as Self::T;
+        const MODULUS_128: u128 = #modulus;
+        const GENERATOR: SmallFp<Self> = SmallFp::new(#generator_mont as Self::T);
+        const ZERO: SmallFp<Self> = SmallFp::new(0 as Self::T);
+        const ONE: SmallFp<Self> = SmallFp::new(#one_mont as Self::T);
+        const NEG_ONE: SmallFp<Self> = SmallFp::new(#neg_one_mont as Self::T);
+
+
+        const TWO_ADICITY: u32 = #two_adicity;
+        const TWO_ADIC_ROOT_OF_UNITY: SmallFp<Self> = SmallFp::new(#two_adic_root_mont as Self::T);
+        #sqrt_precomp_impl
+
+        #[inline(always)]
+        fn add_assign(a: &mut SmallFp<Self>, b: &SmallFp<Self>) {
+            let (mut val, overflow) = a.value.overflowing_add(b.value);
+
+            if overflow {
+                val = Self::T::MAX - Self::MODULUS + 1 + val
+            }
+
+            if val >= Self::MODULUS {
+                val -= Self::MODULUS;
+            }
+            a.value = val;
+        }
+
+        #[inline(always)]
+        fn sub_assign(a: &mut SmallFp<Self>, b: &SmallFp<Self>) {
+            if a.value >= b.value {
+                a.value -= b.value;
+            } else {
+                a.value = Self::MODULUS - (b.value - a.value);
+            }
+        }
+
+        #[inline(always)]
+        fn double_in_place(a: &mut SmallFp<Self>) {
+            let tmp = *a;
+            Self::add_assign(a, &tmp);
+        }
+
+        #[inline(always)]
+        fn neg_in_place(a: &mut SmallFp<Self>) {
+            if a.value != (0 as Self::T) {
+                a.value = Self::MODULUS - a.value;
+            }
+        }
+
+        #mul_impl
+
+        #[inline(always)]
+        fn sum_of_products<const T: usize>(
+            a: &[SmallFp<Self>; T],
+            b: &[SmallFp<Self>; T],) -> SmallFp<Self> {
+            let mut acc = SmallFp::new(0 as Self::T);
+            for (x, y) in a.iter().zip(b.iter()) {
+                let mut prod = *x;
+                Self::mul_assign(&mut prod, y);
+                Self::add_assign(&mut acc, &prod);
+            }
+            acc
+        }
+
+        #[inline(always)]
+        fn square_in_place(a: &mut SmallFp<Self>) {
+            let tmp = *a;
+            Self::mul_assign(a, &tmp);
+        }
+
+        fn inverse(a: &SmallFp<Self>) -> Option<SmallFp<Self>> {
+            if a.value == 0 {
+                return None;
+            }
+
+            let mut result = Self::ONE;
+            let mut base = *a;
+            let mut exp = Self::MODULUS - 2;
+
+            while exp > 0 {
+                if exp & 1 == 1 {
+                    Self::mul_assign(&mut result, &base);
+                }
+
+                let mut sq = base;
+                Self::square_in_place(&mut sq);
+                base = sq;
+                exp >>= 1;
+            }
+
+            Some(result)
+        }
+
+        #from_bigint_impl
+
+        #into_bigint_impl
+    }
+}
+
+// Selects the appropriate multiplication algorithm at compile time:
+// if modulus <= u64, multiply by casting to the next largest primitive
+// otherwise, multiply in parts to form a 256-bit product before reduction
+fn generate_mul_impl(
+    ty: &proc_macro2::TokenStream,
+    modulus: u128,
+    k_bits: u32,
+    r_mask: u128,
+    n_prime: u128,
+) -> proc_macro2::TokenStream {
+    let ty_str = ty.to_string();
+
+    if ty_str == "u128" {
+        quote! {
+            #[inline(always)]
+            fn mul_assign(a: &mut SmallFp<Self>, b: &SmallFp<Self>) {
+                // 256-bit result stored as lo, hi
+                // t = a * b
+                let lolo = (a.value & 0xFFFFFFFFFFFFFFFF) * (b.value & 0xFFFFFFFFFFFFFFFF);
+                let lohi = (a.value & 0xFFFFFFFFFFFFFFFF) * (b.value >> 64);
+                let hilo = (a.value >> 64) * (b.value & 0xFFFFFFFFFFFFFFFF);
+                let hihi = (a.value >> 64) * (b.value >> 64);
+
+                let (cross_sum, cross_carry) = lohi.overflowing_add(hilo);
+                let (mid, mid_carry) = lolo.overflowing_add(cross_sum << 64);
+                let t_lo = mid;
+                let t_hi = hihi + (cross_sum >> 64) + ((cross_carry as u128) << 64) + (mid_carry as u128);
+
+                // m = t_lo * n_prime & r_mask
+                let m = t_lo.wrapping_mul(#n_prime) & #r_mask;
+
+                // mn = m * modulus
+                let lolo = (m & 0xFFFFFFFFFFFFFFFF) * (#modulus & 0xFFFFFFFFFFFFFFFF);
+                let lohi = (m & 0xFFFFFFFFFFFFFFFF) * (#modulus >> 64);
+                let hilo = (m >> 64) * (#modulus & 0xFFFFFFFFFFFFFFFF);
+                let hihi = (m >> 64) * (#modulus >> 64);
+
+                let (cross_sum, cross_carry) = lohi.overflowing_add(hilo);
+                let (mid, mid_carry) = lolo.overflowing_add(cross_sum << 64);
+                let mn_lo = mid;
+                let mn_hi = hihi + (cross_sum >> 64) + ((cross_carry as u128) << 64) + (mid_carry as u128);
+
+                // (t + mn) / R
+                let (sum_lo, carry) = t_lo.overflowing_add(mn_lo);
+                let sum_hi = t_hi + mn_hi + (carry as u128);
+
+                let mut u = (sum_lo >> #k_bits) | (sum_hi << (128 - #k_bits));
+                u -= #modulus * (u >= #modulus) as u128;
+                a.value = u as Self::T;
+            }
+        }
+    } else {
+        let (mul_ty, bits) = match ty_str.as_str() {
+            "u8" => (quote! {u16}, 16u32),
+            "u16" => (quote! {u32}, 32u32),
+            "u32" => (quote! {u64}, 64u32),
+            _ => (quote! {u128}, 128u32),
+        };
+
+        let r_mask_downcast = quote! { #r_mask as #mul_ty };
+        let n_prime_downcast = quote! { #n_prime as #mul_ty };
+        let modulus_downcast = quote! { #modulus as #mul_ty };
+        let one = quote! { 1 as #mul_ty };
+
+        quote! {
+            #[inline(always)]
+            fn mul_assign(a: &mut SmallFp<Self>, b: &SmallFp<Self>) {
+                let a_val = a.value as #mul_ty;
+                let b_val = b.value as #mul_ty;
+
+                let t = a_val * b_val;
+                let t_low = t & #r_mask_downcast;
+
+                // m = t_lo * n_prime & r_mask
+                let m = t_low.wrapping_mul(#n_prime_downcast) & #r_mask_downcast;
+
+                // mn = m * modulus
+                let mn = m * #modulus_downcast;
+
+                // (t + mn) / R
+                let (sum, overflow) = t.overflowing_add(mn);
+                let mut u = sum >> #k_bits;
+
+                u += ((#one) << (#bits - #k_bits)) * (overflow as #mul_ty);
+                u -= #modulus_downcast * ((u >= #modulus_downcast) as #mul_ty);
+                a.value = u as Self::T;
+            }
+        }
+    }
+}
+
+fn mod_inverse_pow2(n: u128, k_bits: u32) -> u128 {
+    let mut inv = 1u128;
+    for _ in 0..k_bits {
+        inv = inv.wrapping_mul(2u128.wrapping_sub(n.wrapping_mul(inv)));
+    }
+    let mask = (1u128 << k_bits) - 1;
+    inv.wrapping_neg() & mask
+}
+
+pub(crate) fn new(modulus: u128, _ty: proc_macro2::TokenStream) -> proc_macro2::TokenStream {
+    let k_bits = 128 - modulus.leading_zeros();
+    let r: u128 = 1u128 << k_bits;
+    let r_mod_n = r % modulus;
+    let r2 = mod_mul_const(r_mod_n, r_mod_n, modulus);
+
+    quote! {
+        pub fn new(value: <Self as SmallFpConfig>::T) -> SmallFp<Self> {
+            let reduced_value = value % <Self as SmallFpConfig>::MODULUS;
+            let mut tmp = SmallFp::new(reduced_value);
+            let r2_elem = SmallFp::new(#r2 as <Self as SmallFpConfig>::T);
+            <Self as SmallFpConfig>::mul_assign(&mut tmp, &r2_elem);
+            tmp
+        }
+
+        pub fn exit(a: &mut SmallFp<Self>) {
+            let mut tmp = *a;
+            let one = SmallFp::new(1 as <Self as SmallFpConfig>::T);
+            <Self as SmallFpConfig>::mul_assign(&mut tmp, &one);
+            a.value = tmp.value;
+        }
+    }
+}