diff --git a/bench-templates/src/macros/field.rs b/bench-templates/src/macros/field.rs index 585b5000a..62063205b 100644 --- a/bench-templates/src/macros/field.rs +++ b/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::>(); + conversions.bench_function("From u64", |b| { + let mut i = 0; + b.iter(|| { + i = (i + 1) % SAMPLES; + <$F>::from_u64(u64s[i]) + }) + }); conversions.finish() } }; diff --git a/ff/src/fields/models/fp/mod.rs b/ff/src/fields/models/fp/mod.rs index 7f417ed9a..be8b3c923 100644 --- a/ff/src/fields/models/fp/mod.rs +++ b/ff/src/fields/models/fp/mod.rs @@ -64,6 +64,10 @@ pub trait FpConfig: Send + Sync + 'static + Sized { /// which works for every modulus. const SQRT_PRECOMP: Option>>; + /// 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; PRECOMP_TABLE_SIZE]; + /// Set a += b. fn add_assign(a: &mut Fp, b: &Fp); @@ -96,6 +100,11 @@ pub trait FpConfig: Send + Sync + 'static + Sized { /// Convert a field element to an integer in the range `0..(Self::MODULUS - /// 1)`. fn into_bigint(other: Fp) -> BigInt; + + /// 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>; } /// Represents an element of the prime field F_p, where `p == P::MODULUS`. @@ -350,6 +359,11 @@ impl, const N: usize> PrimeField for Fp { fn into_bigint(self) -> BigInt { P::into_bigint(self) } + + #[inline] + fn from_u64(r: u64) -> Option { + P::from_u64(r) + } } impl, const N: usize> FftField for Fp { diff --git a/ff/src/fields/models/fp/montgomery_backend.rs b/ff/src/fields/models/fp/montgomery_backend.rs index 76344227c..ac0b52e90 100644 --- a/ff/src/fields/models/fp/montgomery_backend.rs +++ b/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 @@ pub trait MontConfig: 'static + Sync + Send + Sized { const SQRT_PRECOMP: Option, N>>> = sqrt_precomputation::(); + #[allow(long_running_const_eval)] + const SMALL_ELEMENT_MONTGOMERY_PRECOMP: [Fp, N>; PRECOMP_TABLE_SIZE] = + small_element_montgomery_precomputation::(); + /// (MODULUS + 1) / 4 when MODULUS % 4 == 3. Used for square root precomputations. #[doc(hidden)] const MODULUS_PLUS_ONE_DIV_FOUR: Option> = { @@ -354,6 +364,16 @@ pub trait MontConfig: 'static + Sync + Send + Sized { } } + fn from_u64(r: u64) -> Option, N>> { + if r < PRECOMP_TABLE_SIZE as u64 { + Some(Self::SMALL_ELEMENT_MONTGOMERY_PRECOMP[r as usize]) + } else if BigInt::from(r) >= >::MODULUS { + None + } else { + Some(Fp::new_unchecked(Self::R2).mul_u64(r)) + } + } + fn from_bigint(r: BigInt) -> Option, N>> { let mut r = Fp::new_unchecked(r); if r.is_zero() { @@ -559,6 +579,23 @@ pub const fn sqrt_precomputation>( } } +/// 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>( +) -> [Fp, N>; PRECOMP_TABLE_SIZE] { + let mut lookup_table: [Fp, N>; PRECOMP_TABLE_SIZE] = + [, 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] = , N>>::new(BigInt::new(limbs)); + i += 1; + } + lookup_table +} + /// Construct a [`Fp, 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 @@ impl, const N: usize> FpConfig for MontBackend { const SMALL_SUBGROUP_BASE_ADICITY: Option = T::SMALL_SUBGROUP_BASE_ADICITY; const LARGE_SUBGROUP_ROOT_OF_UNITY: Option> = T::LARGE_SUBGROUP_ROOT_OF_UNITY; const SQRT_PRECOMP: Option>> = T::SQRT_PRECOMP; + const SMALL_ELEMENT_MONTGOMERY_PRECOMP: [Fp; PRECOMP_TABLE_SIZE] = + T::SMALL_ELEMENT_MONTGOMERY_PRECOMP; fn add_assign(a: &mut Fp, b: &Fp) { T::add_assign(a, b) @@ -675,6 +714,10 @@ impl, const N: usize> FpConfig for MontBackend { fn into_bigint(a: Fp) -> BigInt { T::into_bigint(a) } + + fn from_u64(r: u64) -> Option> { + T::from_u64(r) + } } impl, const N: usize> Fp, N> { @@ -787,6 +830,37 @@ impl, const N: usize> Fp, N> { } } + #[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 @@ impl, const N: usize> Fp, N> { #[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( diff --git a/ff/src/fields/prime.rs b/ff/src/fields/prime.rs index 8a8fe89b4..7c2e50e0d 100644 --- a/ff/src/fields/prime.rs +++ b/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; + /// 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.