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() % 65536)
+                .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>; 65536];
+
     /// 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

@@ -2,6 +2,7 @@ use ark_std::{marker::PhantomData, Zero};
 
 use super::{Fp, FpConfig};
 use crate::{biginteger::arithmetic as fa, BigInt, BigInteger, PrimeField, SqrtPrecomputation};
+use crate::fields::Field;
 use ark_ff_macros::unroll_for_loops;
 
 /// A trait that specifies the constants and arithmetic procedures
@@ -76,6 +77,10 @@ pub trait MontConfig<const N: usize>: 'static + Sync + Send + Sized {
     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>; 65536] =
+        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 +359,14 @@ pub trait MontConfig<const N: usize>: 'static + Sync + Send + Sized {
         }
     }
 
+    fn from_u64(r: u64) -> Option<Fp<MontBackend<Self, N>, N>> {
+        if r < 65536 {
+            Some(Self::SMALL_ELEMENT_MONTGOMERY_PRECOMP[r as usize])
+        } else {
+            Self::from_bigint(r.into())
+        }
+    }
+
     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 +572,21 @@ pub const fn sqrt_precomputation<const N: usize, T: MontConfig<N>>(
     }
 }
 
+pub const fn small_element_montgomery_precomputation<const N: usize, T: MontConfig<N>>(
+) -> [Fp<MontBackend<T, N>, N>; 65536] {
+    let mut lookup_table: [Fp<MontBackend<T, N>, N>; 65536] =
+        [<Fp<MontBackend<T, N>, N>>::ZERO; 65536];
+
+    let mut i: usize = 1;
+    while i < 65536 {
+        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 +652,8 @@ impl<T: MontConfig<N>, const N: usize> FpConfig<N> for MontBackend<T, N> {
     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>; 65536] =
+        T::SMALL_ELEMENT_MONTGOMERY_PRECOMP;
 
     fn add_assign(a: &mut Fp<Self, N>, b: &Fp<Self, N>) {
         T::add_assign(a, b)
@@ -675,6 +705,10 @@ impl<T: MontConfig<N>, const N: usize> FpConfig<N> for MontBackend<T, N> {
     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> {

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.