Add new field implementation: 32x9 to be used in webnode

o1-labs/proof-systems#2638
openmina · Nov 29, 2024 · f42a8ec · f42a8ec
1 parent 33a1de2
commit f42a8ec
Show file tree

Hide file tree

Showing 18 changed files with 1,662 additions and 105 deletions.
diff --git a/ec/src/models/bw6/mod.rs b/ec/src/models/bw6/mod.rs
@@ -2,11 +2,11 @@ use crate::{
     models::{ModelParameters, SWModelParameters},
     PairingEngine,
 };
-use ark_ff::fields::{
+use ark_ff::{fields::{
     fp3::Fp3Parameters,
     fp6_2over3::{Fp6, Fp6Parameters},
     BitIteratorBE, Field, PrimeField, SquareRootField,
-};
+}, BigInteger};
 use num_traits::One;
 
 use core::marker::PhantomData;
@@ -68,7 +68,7 @@ impl<P: BW6Parameters> BW6<P> {
     }
 
     fn exp_by_x(mut f: Fp6<P::Fp6Params>) -> Fp6<P::Fp6Params> {
-        f = f.cyclotomic_exp(&P::X);
+        f = f.cyclotomic_exp(&P::X.to_64x4());
         if P::X_IS_NEGATIVE {
             f.conjugate();
         }

diff --git a/ec/src/models/mnt4/mod.rs b/ec/src/models/mnt4/mod.rs
@@ -1,3 +1,5 @@
+use ark_ff::BigInteger;
+
 use {
     crate::{
         models::{ModelParameters, SWModelParameters},
@@ -179,11 +181,11 @@ impl<P: MNT4Parameters> MNT4<P> {
         let mut elt_q = *elt;
         elt_q.frobenius_map(1);
 
-        let w1_part = elt_q.cyclotomic_exp(&P::FINAL_EXPONENT_LAST_CHUNK_1);
+        let w1_part = elt_q.cyclotomic_exp(&P::FINAL_EXPONENT_LAST_CHUNK_1.to_64x4());
         let w0_part = if P::FINAL_EXPONENT_LAST_CHUNK_W0_IS_NEG {
-            elt_inv_clone.cyclotomic_exp(&P::FINAL_EXPONENT_LAST_CHUNK_ABS_OF_W0)
+            elt_inv_clone.cyclotomic_exp(&P::FINAL_EXPONENT_LAST_CHUNK_ABS_OF_W0.to_64x4())
         } else {
-            elt_clone.cyclotomic_exp(&P::FINAL_EXPONENT_LAST_CHUNK_ABS_OF_W0)
+            elt_clone.cyclotomic_exp(&P::FINAL_EXPONENT_LAST_CHUNK_ABS_OF_W0.to_64x4())
         };
 
         w1_part * &w0_part

diff --git a/ec/src/models/mnt6/mod.rs b/ec/src/models/mnt6/mod.rs
@@ -1,3 +1,5 @@
+use ark_ff::BigInteger;
+
 use {
     crate::{
         models::{ModelParameters, SWModelParameters},
@@ -185,11 +187,11 @@ impl<P: MNT6Parameters> MNT6<P> {
         let mut elt_q = *elt;
         elt_q.frobenius_map(1);
 
-        let w1_part = elt_q.cyclotomic_exp(&P::FINAL_EXPONENT_LAST_CHUNK_1);
+        let w1_part = elt_q.cyclotomic_exp(&P::FINAL_EXPONENT_LAST_CHUNK_1.to_64x4());
         let w0_part = if P::FINAL_EXPONENT_LAST_CHUNK_W0_IS_NEG {
-            elt_inv_clone.cyclotomic_exp(&P::FINAL_EXPONENT_LAST_CHUNK_ABS_OF_W0)
+            elt_inv_clone.cyclotomic_exp(&P::FINAL_EXPONENT_LAST_CHUNK_ABS_OF_W0.to_64x4())
         } else {
-            elt_clone.cyclotomic_exp(&P::FINAL_EXPONENT_LAST_CHUNK_ABS_OF_W0)
+            elt_clone.cyclotomic_exp(&P::FINAL_EXPONENT_LAST_CHUNK_ABS_OF_W0.to_64x4())
         };
 
         w1_part * &w0_part

diff --git a/ec/src/models/short_weierstrass_jacobian.rs b/ec/src/models/short_weierstrass_jacobian.rs
@@ -11,9 +11,7 @@ use ark_std::{
 };
 
 use ark_ff::{
-    bytes::{FromBytes, ToBytes},
-    fields::{BitIteratorBE, Field, PrimeField, SquareRootField},
-    ToConstraintField, UniformRand,
+    bytes::{FromBytes, ToBytes}, fields::{BitIteratorBE, Field, PrimeField, SquareRootField}, BigInteger, ToConstraintField, UniformRand
 };
 
 use crate::{models::SWModelParameters as Parameters, AffineCurve, ProjectiveCurve};
@@ -223,7 +221,8 @@ impl<P: Parameters> AffineCurve for GroupAffine<P> {
 
     #[inline]
     fn mul<S: Into<<Self::ScalarField as PrimeField>::BigInt>>(&self, by: S) -> GroupProjective<P> {
-        let bits = BitIteratorBE::new(by.into());
+        let inner: <Self::ScalarField as PrimeField>::BigInt = by.into();
+        let bits = ark_ff::BitIteratorBE::new(inner.to_64x4());
         self.mul_bits(bits)
     }
 
@@ -714,7 +713,7 @@ impl<'a, P: Parameters> SubAssign<&'a Self> for GroupProjective<P> {
 
 impl<P: Parameters> MulAssign<P::ScalarField> for GroupProjective<P> {
     fn mul_assign(&mut self, other: P::ScalarField) {
-        *self = self.mul(other.into_repr())
+        *self = self.mul(other.into_repr().to_64x4())
     }
 }
 

diff --git a/ec/src/models/twisted_edwards_extended.rs b/ec/src/models/twisted_edwards_extended.rs
@@ -21,9 +21,7 @@ use num_traits::{One, Zero};
 use zeroize::Zeroize;
 
 use ark_ff::{
-    bytes::{FromBytes, ToBytes},
-    fields::{BitIteratorBE, Field, PrimeField, SquareRootField},
-    ToConstraintField, UniformRand,
+    bytes::{FromBytes, ToBytes}, fields::{BitIteratorBE, Field, PrimeField, SquareRootField}, BigInteger, ToConstraintField, UniformRand
 };
 
 #[cfg(feature = "parallel")]
@@ -138,7 +136,8 @@ impl<P: Parameters> AffineCurve for GroupAffine<P> {
     }
 
     fn mul<S: Into<<Self::ScalarField as PrimeField>::BigInt>>(&self, by: S) -> GroupProjective<P> {
-        self.mul_bits(BitIteratorBE::new(by.into()))
+        let inner: <Self::ScalarField as PrimeField>::BigInt = by.into();
+        self.mul_bits(ark_ff::BitIteratorBE::new(inner.to_64x4()))
     }
 
     fn from_random_bytes(bytes: &[u8]) -> Option<Self> {
@@ -610,7 +609,7 @@ impl<'a, P: Parameters> SubAssign<&'a Self> for GroupProjective<P> {
 
 impl<P: Parameters> MulAssign<P::ScalarField> for GroupProjective<P> {
     fn mul_assign(&mut self, other: P::ScalarField) {
-        *self = self.mul(other.into_repr())
+        *self = self.mul(other.into_repr().to_64x4())
     }
 }
 

diff --git a/ec/src/msm/variable_base.rs b/ec/src/msm/variable_base.rs
@@ -54,7 +54,7 @@ impl VariableBaseMSM {
                         scalar.divn(w_start as u32);
 
                         // We mod the remaining bits by 2^{window size}, thus taking `c` bits.
-                        let scalar = scalar.as_ref()[0] % (1 << c);
+                        let scalar = scalar.to_64x4()[0] % (1 << c);
 
                         // If the scalar is non-zero, we update the corresponding
                         // bucket.

diff --git a/ff/Cargo.toml b/ff/Cargo.toml
@@ -33,3 +33,7 @@ default = []
 std = [ "ark-std/std", "ark-serialize/std" ]
 parallel = [ "std", "rayon", "ark-std/parallel" ]
 asm = []
+32x9 = []
+
+[lints.rust]
+unexpected_cfgs = { level = "warn", check-cfg = ['cfg(use_asm)'] }
diff --git a/ff/src/biginteger/arithmetic.rs b/ff/src/biginteger/arithmetic.rs
@@ -1,3 +1,4 @@
+#![allow(unused)]
 use ark_std::vec::Vec;
 
 /// Make 4 u64 multiplications, instead of 1 u128

diff --git a/ff/src/biginteger/macros.rs b/ff/src/biginteger/macros.rs
@@ -1,17 +1,43 @@
 macro_rules! bigint_impl {
     ($name:ident, $num_limbs:expr) => {
         #[derive(Copy, Clone, PartialEq, Eq, Debug, Default, Hash, Zeroize)]
-        pub struct $name(pub [u64; $num_limbs]);
+        pub struct $name(pub(crate) [u64; $num_limbs]);
 
         impl $name {
             pub const fn new(value: [u64; $num_limbs]) -> Self {
                 $name(value)
             }
+
+            pub const fn to_64x4(&self) -> [u64; $num_limbs] {
+                self.0
+            }
+
+            pub const fn from_64x4(value: [u64; $num_limbs]) -> Self {
+                $name(value)
+            }
+
+            #[ark_ff_asm::unroll_for_loops]
+            pub fn assign_bits_and(&mut self, other: &Self) {
+                for i in 0..$num_limbs {
+                    self.0[i] |= other.0[i]
+                }
+            }
+
+            pub fn to_native(&self) -> [u64; $num_limbs] {
+                self.0
+            }
         }
 
         impl BigInteger for $name {
             const NUM_LIMBS: usize = $num_limbs;
 
+            fn to_64x4(&self) -> [u64; 4] {
+                self.0
+            }
+            fn from_64x4(value: [u64; 4]) -> Self {
+                $name(value)
+            }
+
             #[inline]
             #[ark_ff_asm::unroll_for_loops]
             fn add_nocarry(&mut self, other: &Self) -> bool {

diff --git a/ff/src/biginteger/mod.rs b/ff/src/biginteger/mod.rs
@@ -1,7 +1,5 @@
 use crate::{
-    bytes::{FromBytes, ToBytes},
-    fields::{BitIteratorBE, BitIteratorLE},
-    UniformRand,
+    bytes::{FromBytes, ToBytes}, fields::{BitIteratorBE, BitIteratorLE}, UniformRand
 };
 use ark_serialize::{CanonicalDeserialize, CanonicalSerialize, SerializationError};
 use ark_std::rand::{
@@ -31,14 +29,17 @@ pub fn signed_mod_reduction(n: u64, modulus: u64) -> i64 {
     }
 }
 
-bigint_impl!(BigInteger64, 1);
-bigint_impl!(BigInteger128, 2);
-bigint_impl!(BigInteger256, 4);
-bigint_impl!(BigInteger320, 5);
-bigint_impl!(BigInteger384, 6);
-bigint_impl!(BigInteger448, 7);
-bigint_impl!(BigInteger768, 12);
-bigint_impl!(BigInteger832, 13);
+pub mod native_bigint {
+    use super::*;
+    bigint_impl!(BigInteger256, 4);
+}
+pub mod webnode;
+
+#[cfg(not(any(target_family = "wasm", feature = "32x9")))]
+pub use native_bigint::*;
+
+#[cfg(any(target_family = "wasm", feature = "32x9"))]
+pub use webnode::*;
 
 #[cfg(test)]
 mod tests;
@@ -63,15 +64,16 @@ pub trait BigInteger:
     + 'static
     + UniformRand
     + Zeroize
-    + AsMut<[u64]>
-    + AsRef<[u64]>
     + From<u64>
     + TryFrom<BigUint>
     + Into<BigUint>
 {
     /// Number of limbs.
     const NUM_LIMBS: usize;
 
+    fn to_64x4(&self) -> [u64; 4];
+    fn from_64x4(value: [u64; 4]) -> Self;
+
     /// Add another representation to this one, returning the carry bit.
     fn add_nocarry(&mut self, other: &Self) -> bool;
 
@@ -119,13 +121,13 @@ pub trait BigInteger:
     /// Returns the bit representation in a big endian boolean array,
     /// with leading zeroes.
     fn to_bits_be(&self) -> Vec<bool> {
-        BitIteratorBE::new(self).collect::<Vec<_>>()
+        BitIteratorBE::new(self.to_64x4()).collect::<Vec<_>>()
     }
 
     /// Returns the bit representation in a little endian boolean array,
     /// with trailing zeroes.
     fn to_bits_le(&self) -> Vec<bool> {
-        BitIteratorLE::new(self).collect::<Vec<_>>()
+        BitIteratorLE::new(self.to_64x4()).collect::<Vec<_>>()
     }
 
     /// Returns the byte representation in a big endian byte array,
@@ -143,11 +145,12 @@ pub trait BigInteger:
         if w >= 2 && w < 64 {
             let mut res = vec![];
             let mut e = *self;
+            let e64 = self.to_64x4();
 
             while !e.is_zero() {
                 let z: i64;
                 if e.is_odd() {
-                    z = signed_mod_reduction(e.as_ref()[0], 1 << w);
+                    z = signed_mod_reduction(e64.as_ref()[0], 1 << w);
                     if z >= 0 {
                         e.sub_noborrow(&Self::from(z as u64));
                     } else {