[do not merge] Introduce optimized from_u64 function for conversion to Montgomery form

bench-templates/src/macros/field.rs

@@ -6,6 +6,7 @@ macro_rules! f_bench {
             mod [<$F:lower>] {
                 use super::*;
                 use ark_ff::{Field, PrimeField, UniformRand};
+                use ark_std::rand::RngCore;
                 field_common!($bench_group_name, $F);
                 sqrt!($bench_group_name, $F);
                 prime_field!($bench_group_name, $F);
@@ -402,6 +403,16 @@ macro_rules! prime_field {
                     f[i].into_bigint()
                 })
             });
+            let u64s = (0..SAMPLES)
+                .map(|_| rng.next_u64())
+                .collect::<Vec<_>>();
+            conversions.bench_function("From u64", |b| {
+                let mut i = 0;
+                b.iter(|| {
+                    i = (i + 1) % SAMPLES;
+                    <$F>::from_u64(u64s[i])
+                })
+            });            
             conversions.finish()
         }
     };

ff/src/fields/models/fp/mod.rs

@@ -64,6 +64,10 @@ pub trait FpConfig<const N: usize>: Send + Sync + 'static + Sized {
     /// which works for every modulus.
     const SQRT_PRECOMP: Option<SqrtPrecomputation<Fp<Self, N>>>;
 
+    /// Precomputed lookup table for values 0..2^16 in Montgomery form.
+    /// Otherwise, conversion to Montgomery form requires a multiplication by R^2.
+    const SMALL_ELEMENT_MONTGOMERY_PRECOMP: [Fp<Self, N>; PRECOMP_TABLE_SIZE];
+
     /// Set a += b.
     fn add_assign(a: &mut Fp<Self, N>, b: &Fp<Self, N>);
 
@@ -96,6 +100,11 @@ pub trait FpConfig<const N: usize>: Send + Sync + 'static + Sized {
     /// Convert a field element to an integer in the range `0..(Self::MODULUS -
     /// 1)`.
     fn into_bigint(other: Fp<Self, N>) -> BigInt<N>;
+
+    /// Construct a field element from a u64 in the range
+    /// `0..(Self::MODULUS - 1)`. Returns `None` if the integer is outside
+    /// this range.
+    fn from_u64(other: u64) -> Option<Fp<Self, N>>;    
 }
 
 /// Represents an element of the prime field F_p, where `p == P::MODULUS`.
@@ -350,6 +359,11 @@ impl<P: FpConfig<N>, const N: usize> PrimeField for Fp<P, N> {
     fn into_bigint(self) -> BigInt<N> {
         P::into_bigint(self)
     }
+
+    #[inline]
+    fn from_u64(r: u64) -> Option<Self> {
+        P::from_u64(r)
+    }    
 }
 
 impl<P: FpConfig<N>, const N: usize> FftField for Fp<P, N> {

ff/src/fields/models/fp/montgomery_backend.rs

@@ -1,9 +1,15 @@
 use ark_std::{marker::PhantomData, Zero};
 
 use super::{Fp, FpConfig};
-use crate::{biginteger::arithmetic as fa, BigInt, BigInteger, PrimeField, SqrtPrecomputation};
+use crate::{
+    biginteger::arithmetic::{self as fa},
+    fields::Field,
+    BigInt, BigInteger, PrimeField, SqrtPrecomputation,
+};
 use ark_ff_macros::unroll_for_loops;
 
+pub const PRECOMP_TABLE_SIZE: usize = 1 << 14;
+
 /// A trait that specifies the constants and arithmetic procedures
 /// for Montgomery arithmetic over the prime field defined by `MODULUS`.
 ///
@@ -76,6 +82,10 @@
     const SQRT_PRECOMP: Option<SqrtPrecomputation<Fp<MontBackend<Self, N>, N>>> =
         sqrt_precomputation::<N, Self>();
 
+    #[allow(long_running_const_eval)]
+    const SMALL_ELEMENT_MONTGOMERY_PRECOMP: [Fp<MontBackend<Self, N>, N>; PRECOMP_TABLE_SIZE] =
+        small_element_montgomery_precomputation::<N, Self>();
+
     /// (MODULUS + 1) / 4 when MODULUS % 4 == 3. Used for square root precomputations.
     #[doc(hidden)]
     const MODULUS_PLUS_ONE_DIV_FOUR: Option<BigInt<N>> = {
@@ -354,6 +364,16 @@
         }
     }
 
+    fn from_u64(r: u64) -> Option<Fp<MontBackend<Self, N>, N>> {
+        if r < PRECOMP_TABLE_SIZE as u64 {
+            Some(Self::SMALL_ELEMENT_MONTGOMERY_PRECOMP[r as usize])
+        } else if BigInt::from(r) >= <MontBackend<Self, N>>::MODULUS {
+            None
+        } else {
+            Some(Fp::new_unchecked(Self::R2).mul_u64(r))
+        }
+    }
+
     fn from_bigint(r: BigInt<N>) -> Option<Fp<MontBackend<Self, N>, N>> {
         let mut r = Fp::new_unchecked(r);
         if r.is_zero() {
@@ -559,6 +579,23 @@
     }
 }
 
+/// Adapted the `bn256-table` feature from `halo2curves`:
+/// https://github.com/privacy-scaling-explorations/halo2curves/blob/main/script/bn256.py
+pub const fn small_element_montgomery_precomputation<const N: usize, T: MontConfig<N>>(
+) -> [Fp<MontBackend<T, N>, N>; PRECOMP_TABLE_SIZE] {
+    let mut lookup_table: [Fp<MontBackend<T, N>, N>; PRECOMP_TABLE_SIZE] =
+        [<Fp<MontBackend<T, N>, N>>::ZERO; PRECOMP_TABLE_SIZE];
+
+    let mut i: usize = 1;
+    while i < PRECOMP_TABLE_SIZE {
+        let mut limbs = [0u64; N];
+        limbs[0] = i as u64;
+        lookup_table[i] = <Fp<MontBackend<T, N>, N>>::new(BigInt::new(limbs));
+        i += 1;
+    }
+    lookup_table
+}
+
 /// Construct a [`Fp<MontBackend<T, N>, N>`] element from a literal string. This
 /// should be used primarily for constructing constant field elements; in a
 /// non-const context, [`Fp::from_str`](`ark_std::str::FromStr::from_str`) is
@@ -624,6 +661,8 @@
     const SMALL_SUBGROUP_BASE_ADICITY: Option<u32> = T::SMALL_SUBGROUP_BASE_ADICITY;
     const LARGE_SUBGROUP_ROOT_OF_UNITY: Option<Fp<Self, N>> = T::LARGE_SUBGROUP_ROOT_OF_UNITY;
     const SQRT_PRECOMP: Option<crate::SqrtPrecomputation<Fp<Self, N>>> = T::SQRT_PRECOMP;
+    const SMALL_ELEMENT_MONTGOMERY_PRECOMP: [Fp<Self, N>; PRECOMP_TABLE_SIZE] =
+        T::SMALL_ELEMENT_MONTGOMERY_PRECOMP;
 
     fn add_assign(a: &mut Fp<Self, N>, b: &Fp<Self, N>) {
         T::add_assign(a, b)
@@ -675,6 +714,10 @@
     fn into_bigint(a: Fp<Self, N>) -> BigInt<N> {
         T::into_bigint(a)
     }
+
+    fn from_u64(r: u64) -> Option<Fp<Self, N>> {
+        T::from_u64(r)
+    }
 }
 
 impl<T: MontConfig<N>, const N: usize> Fp<MontBackend<T, N>, N> {
@@ -787,6 +830,37 @@
         }
     }
 
+    #[unroll_for_loops(12)]
+    #[inline(always)]
+    pub fn mul_u64(mut self, other: u64) -> Self {
+        let (mut lo, mut hi) = ([0u64; N], [0u64; N]);
+
+        for i in 0..N - 1 {
+            lo[i] = mac_with_carry!(lo[i], (self.0).0[i], other, &mut lo[i + 1]);
+        }
+        lo[N - 1] = mac_with_carry!(lo[N - 1], (self.0).0[N - 1], other, &mut hi[0]);
+
+        // Montgomery reduction
+        let mut carry2 = 0;
+        for i in 0..N {
+            let tmp = lo[i].wrapping_mul(T::INV);
+            let mut carry = 0u64;
+            mac!(lo[i], tmp, T::MODULUS.0[0], &mut carry);
+            for j in 1..N {
+                let k: usize = i + j;
+                if k >= N {
+                    hi[k - N] = mac_with_carry!(hi[k - N], tmp, T::MODULUS.0[j], &mut carry);
+                } else {
+                    lo[k] = mac_with_carry!(lo[k], tmp, T::MODULUS.0[j], &mut carry);
+                }
+            }
+            hi[i] = adc!(hi[i], carry, &mut carry2);
+        }
+
+(self.0).0 = hi;
+        self.const_subtract_modulus_with_carry(carry2 != 0)
+    }
+
     const fn const_is_valid(&self) -> bool {
         crate::const_for!((i in 0..N) {
             if (self.0).0[N - i - 1] < T::MODULUS.0[N - i - 1] {
@@ -821,10 +895,21 @@
 
 #[cfg(test)]
 mod test {
-    use ark_std::{str::FromStr, vec::Vec};
+    use ark_std::{rand::RngCore, str::FromStr};
     use ark_test_curves::secp256k1::Fr;
     use num_bigint::{BigInt, BigUint, Sign};
 
+    #[test]
+    fn test_mul_u64() {
+        let r2 = Fr::new_unchecked(Fr::R2);
+        let mut rng = ark_std::test_rng();
+        let value = rng.next_u64();
+        assert_eq!(
+            r2.mul_u64(value),
+            r2 * Fr::new_unchecked(value.into())
+        );
+    }
+
     #[test]
     fn test_mont_macro_correctness() {
         let (is_positive, limbs) = str_to_limbs_u64(

ff/src/fields/prime.rs

@@ -57,6 +57,9 @@ pub trait PrimeField:
     /// Converts an element of the prime field into an integer in the range 0..(p - 1).
     fn into_bigint(self) -> Self::BigInt;
 
+    /// Construct a prime field element from a u64 in the range 0..(p - 1).
+    fn from_u64(repr: u64) -> Option<Self>;
+
     /// Reads bytes in big-endian, and converts them to a field element.
     /// If the integer represented by `bytes` is larger than the modulus `p`, this method
     /// performs the appropriate reduction.