Const initialization helper for SmallFp

CHANGELOG.md

@@ -21,6 +21,7 @@
 ### Features
 
 - (`ark-serialize`) Implementation of `CanonicalSerialize` and `CanonicalDeserialize` for signed integer types
+- [\#1084](https://github.com/arkworks-rs/algebra/pull/1084) (`ark-ff`) Add `from_u128` const constructor for `SmallFp` fields
 
 ### Improvements
 

ff-macros/src/lib.rs

@@ -12,7 +12,7 @@ use proc_macro::TokenStream;
 use quote::format_ident;
 use syn::{Expr, ExprLit, Item, ItemFn, Lit, Meta};
 
-mod montgomery;
+pub(crate) mod montgomery;
 mod small_fp;
 mod unroll;
 
@@ -46,31 +46,50 @@ pub fn define_field(input: TokenStream) -> TokenStream {
 
     let name = args.name;
     let config_name = format_ident!("{}Config", name);
-    let modulus = syn::LitStr::new(&args.modulus, proc_macro2::Span::call_site());
-    let generator = syn::LitStr::new(&args.generator, proc_macro2::Span::call_site());
-
     let is_small_modulus = modulus_big < (BigUint::from(1u128) << 127);
-    let derives = if is_small_modulus {
-        quote::quote! { #[derive(ark_ff::SmallFpConfig)] }
-    } else {
-        quote::quote! { #[derive(ark_ff::MontConfig)] }
-    };
 
-    let field_ty = if is_small_modulus {
-        quote::quote! { ark_ff::SmallFp<#config_name> }
+    if is_small_modulus {
+        let modulus_u128: u128 = args
+            .modulus
+            .parse()
+            .expect("modulus should fit in u128 for small field");
+        let generator_u128: u128 = args
+            .generator
+            .parse()
+            .expect("generator should fit in u128 for small field");
+
+        let config_impl =
+            small_fp::small_fp_config_helper(modulus_u128, generator_u128, config_name.clone());
+
+        quote::quote! {
+            pub struct #config_name;
+            pub type #name = ark_ff::SmallFp<#config_name>;
+            #config_impl
+        }
+        .into()
     } else {
         let fp_alias = utils::fp_alias_for_limbs(limbs);
-        quote::quote! { #fp_alias<ark_ff::MontBackend<#config_name, #limbs>> }
-    };
 
-    quote::quote! {
-        #derives
-        #[modulus = #modulus]
-        #[generator = #generator]
-        pub struct #config_name;
-        pub type #name = #field_ty;
+        let generator_big = args
+            .generator
+            .parse::<BigUint>()
+            .expect("generator should be a decimal integer string");
+
+        let config_impl = montgomery::mont_config_helper(
+            modulus_big,
+            generator_big,
+            None,
+            None,
+            config_name.clone(),
+        );
+
+        quote::quote! {
+            pub struct #config_name;
+            pub type #name = #fp_alias<ark_ff::MontBackend<#config_name, #limbs>>;
+            #config_impl
+        }
+        .into()
     }
-    .into()
 }
 
 /// Derive the `MontConfig` trait.

ff-macros/src/small_fp/mod.rs

@@ -22,8 +22,8 @@ pub(crate) fn small_fp_config_helper(
         "SmallFpConfig montgomery backend supports only moduli < 2^127. Use MontConfig with BigInt instead of SmallFp."
     );
 
-    let backend_impl = montgomery_backend::backend_impl(&ty, modulus, generator);
-    let exit_impl = montgomery_backend::exit_impl();
+    let (backend_impl, r_mod_p) = montgomery_backend::backend_impl(&ty, modulus, generator);
+    let exit_impl = montgomery_backend::exit_impl(modulus, r_mod_p);
 
     quote! {
         const _: () = {

ff-macros/src/small_fp/montgomery_backend.rs

@@ -8,7 +8,7 @@ pub(crate) fn backend_impl(
     ty: &proc_macro2::TokenStream,
     modulus: u128,
     generator: u128,
-) -> proc_macro2::TokenStream {
+) -> (proc_macro2::TokenStream, u128) {
     assert!(modulus > 1, "modulus must be greater than 1");
     assert!(
         modulus % 2 == 1,
@@ -88,7 +88,7 @@ pub(crate) fn backend_impl(
             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
+                    val += Self::T::MAX - Self::MODULUS + 1
                 }
                 if val >= Self::MODULUS {
                     val -= Self::MODULUS;
@@ -98,7 +98,7 @@ pub(crate) fn backend_impl(
         }
     };
 
-    quote! {
+    let ts = quote! {
         type T = #ty;
         const MODULUS: Self::T = #modulus as Self::T;
         const MODULUS_U128: u128 = #modulus;
@@ -209,7 +209,9 @@ pub(crate) fn backend_impl(
         #from_bigint_impl
 
         #into_bigint_impl
-    }
+    };
+
+    (ts, r_mod_n)
 }
 
 // Selects the appropriate multiplication algorithm at compile time:
@@ -489,11 +491,36 @@ fn mod_inverse_pow2(n: u128, k_bits: u32) -> u128 {
     inv.wrapping_neg() & mask
 }
 
-pub(crate) fn exit_impl() -> proc_macro2::TokenStream {
+pub(crate) fn exit_impl(modulus: u128, r_mod_p: u128) -> proc_macro2::TokenStream {
     quote! {
         pub fn exit(a: &mut SmallFp<Self>) {
             let one = SmallFp::from_raw(1 as <Self as SmallFpConfig>::T);
             <Self as SmallFpConfig>::mul_assign(a, &one);
         }
+
+        // This is the `SmallFp` analogue of [`MontFp!`] to support const initialization
+        pub const fn from_u128(value: u128) -> SmallFp<Self> {
+            const MODULUS: u128 = #modulus;
+            const R_MOD_P: u128 = #r_mod_p;
+
+            // const-compatible modular multiplication via double-and-add
+            // Safe from overflow: modulus < 2^127 so a,result < 2^127 and all additions fit u128
+            const fn mod_mul(mut a: u128, mut b: u128, m: u128) -> u128 {
+                a %= m;
+                let mut result = 0u128;
+                while b > 0 {
+                    if b & 1 == 1 {
+                        result = (result + a) % m;
+                    }
+                    a = (a + a) % m;
+                    b >>= 1;
+                }
+                result
+            }
+
+            let val = value % MODULUS;
+            let mont = mod_mul(val, R_MOD_P, MODULUS);
+            SmallFp::from_raw(mont as <Self as SmallFpConfig>::T)
+        }
     }
 }

ff-macros/src/small_fp/utils.rs

@@ -85,7 +85,7 @@ pub(crate) fn generate_montgomery_bigint_casts(
         quote! {
             fn from_bigint(a: ark_ff::BigInt<2>) -> Option<SmallFp<Self>> {
                 let val = (a.0[0] as u128) + ((a.0[1] as u128) << 64);
-                if val > Self::MODULUS_U128 {
+                if val >= Self::MODULUS_U128 {
                     None
                 } else {
                     let reduced_val = val % Self::MODULUS_U128;

test-curves/src/smallfp.rs

@@ -43,4 +43,83 @@ mod tests {
     test_small_field!(f32_mont_babybear; SmallFp32Babybear);
     test_small_field!(f64; SmallFp64Goldilock);
     test_small_field!(f128; SmallFp128);
+
+    mod const_constructors {
+        use super::*;
+        use ark_ff::{One, Zero};
+
+        #[test]
+        fn test_from_u128_zero() {
+            let a: SmallFp64Goldilock = SmallFp64GoldilockConfig::from_u128(0);
+            assert!(a.is_zero(), "from_u128(0) should be zero");
+        }
+
+        #[test]
+        fn test_from_u128_one() {
+            let a: SmallFp64Goldilock = SmallFp64GoldilockConfig::from_u128(1);
+            assert!(a.is_one(), "from_u128(1) should be one");
+        }
+
+        #[test]
+        fn test_from_u128_matches_runtime() {
+            for val in [
+                0u128,
+                1,
+                2,
+                7,
+                42,
+                255,
+                65521,
+                2013265921,
+                18446744069414584320,
+            ] {
+                let const_elem: SmallFp64Goldilock = SmallFp64GoldilockConfig::from_u128(val);
+                let runtime_elem = SmallFp64Goldilock::from(val);
+                assert_eq!(const_elem, runtime_elem, "from_u128({val}) mismatch");
+            }
+        }
+
+        #[test]
+        fn test_from_u128_all_backing_types() {
+            // u8 field
+            for val in [0u128, 1, 7, 42, 250] {
+                assert_eq!(SmallFp8Config::from_u128(val), SmallFp8::from(val));
+            }
+            // u32 field (M31)
+            for val in [0u128, 1, 7, 1000000, 2147483646] {
+                assert_eq!(SmallFp32M31Config::from_u128(val), SmallFp32M31::from(val));
+            }
+            // u128 field
+            for val in [
+                0u128,
+                1,
+                7,
+                u64::MAX as u128,
+                143244528689204659050391023439224324688,
+            ] {
+                assert_eq!(SmallFp128Config::from_u128(val), SmallFp128::from(val));
+            }
+        }
+
+        #[test]
+        fn test_from_u128_reduction() {
+            let modulus = 18446744069414584321u128; // Goldilocks
+            let val = modulus + 7;
+            let const_elem: SmallFp64Goldilock = SmallFp64GoldilockConfig::from_u128(val);
+            let seven: SmallFp64Goldilock = SmallFp64GoldilockConfig::from_u128(7);
+            assert_eq!(
+                const_elem, seven,
+                "from_u128(P+7) should equal from_u128(7)"
+            );
+        }
+
+        #[test]
+        fn test_const_context() {
+            const SEVEN: SmallFp64Goldilock = SmallFp64GoldilockConfig::from_u128(7);
+            const FORTY_TWO: SmallFp128 = SmallFp128Config::from_u128(42);
+
+            assert_eq!(SEVEN, SmallFp64Goldilock::from(7u128));
+            assert_eq!(FORTY_TWO, SmallFp128::from(42u128));
+        }
+    }
 }