Implement u32 backend for Fp

penumbra-zone · Jan 12, 2024 · c9908e8 · c9908e8
1 parent 031a44a
commit c9908e8
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 27 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -39,6 +39,8 @@ parallel = ["ark-ff/parallel", "ark-ec/parallel", "ark-groth16/parallel", "ark-s
 r1cs = ["ark-r1cs-std", "ark-groth16", "ark-snark"]
 # TODO: eventually, feature-gate all arkworks deps behind this feature.
 arkworks = []
+u32 = []
+u64 = []
 
 [dev-dependencies]
 proptest = "1"

diff --git a/src/fields/fp/u32/arkworks_constants.rs b/src/fields/fp/u32/arkworks_constants.rs
@@ -1,27 +1,41 @@
 use super::Fp;
 
-pub const NUM_LIMBS: usize = 12;
-
-pub const MODULUS_LIMBS: [u32; 12] = [
-    1, 2231943168, 805306368, 386620740, 3121170432, 519266863, 16061327, 438491635, 1822509371,
-    3325756864, 398790890, 28195398,
+pub const MODULUS_LIMBS: [u64; 6] = [
+    9586122913090633729,
+    1660523435060625408,
+    2230234197602682880,
+    1883307231910630287,
+    14284016967150029115,
+    121098312706494698,
 ];
 
-pub const MODULUS_MINUS_ONE_DIV_TWO_LIMBS: [u32; 12] = [
-    0, 1115971584, 402653184, 193310370, 3708068864, 2407117079, 2155514311, 2366729465, 911254685,
-    1662878432, 199395445, 14097699,
+pub const MODULUS_MINUS_ONE_DIV_TWO_LIMBS: [u64; 6] = [
+    4793061456545316864,
+    830261717530312704,
+    10338489135656117248,
+    10165025652810090951,
+    7142008483575014557,
+    60549156353247349,
 ];
 
 pub const MODULUS_BIT_SIZE: u32 = 0x179;
 
-pub const TRACE_LIMBS: [u32; 12] = [
-    136227, 1964032000, 536894509, 2294212645, 1312586701, 1741423572, 619473035, 385987205,
-    1135286508, 3910688532, 1720, 0,
+pub const TRACE_LIMBS: [u64; 6] = [
+    8435453208297608227,
+    9853568280881552429,
+    7479357291536088013,
+    1657802422768920715,
+    16796279350917535980,
+    1720,
 ];
 
-pub const TRACE_MINUS_ONE_DIV_TWO_LIMBS: [u32; 12] = [
-    68113, 3129499648, 2415930902, 3294589970, 656293350, 3018195434, 2457220165, 192993602,
-    567643254, 1955344266, 860, 0,
+pub const TRACE_MINUS_ONE_DIV_TWO_LIMBS: [u64; 6] = [
+    13441098641003579921,
+    14150156177295552022,
+    12963050682622819814,
+    828901211384460357,
+    8398139675458767990,
+    860,
 ];
 
 pub const TWO_ADICITY: u32 = 0x2e;

diff --git a/src/fields/fp/u32/wrapper.rs b/src/fields/fp/u32/wrapper.rs
@@ -5,7 +5,7 @@ pub struct Fp(pub fiat::FpMontgomeryDomainFieldElement);
 
 impl PartialEq for Fp {
     fn eq(&self, other: &Self) -> bool {
-        let sub = self.sub(*other);
+        let sub = self.sub(other);
         let mut check_word = 0;
         fiat::fp_nonzero(&mut check_word, &sub.0 .0);
         check_word == 0
@@ -21,7 +21,15 @@ impl zeroize::Zeroize for Fp {
 }
 
 impl Fp {
-    pub fn from_le_limbs(limbs: [u32; 12]) -> Fp {
+    pub fn from_le_limbs(limbs: [u64; 6]) -> Fp {
+        let limbs = {
+            let mut out = [0u32; 12];
+            for i in 0..6 {
+                out[2 * i] = (limbs[i] & 0xFFFF_FFFF_FFFF_FFFF) as u32;
+                out[2 * i + 1] = (limbs[i] >> 32) as u32;
+            }
+            out
+        };
         let x_non_monty = fiat::FpNonMontgomeryDomainFieldElement(limbs);
         let mut x = fiat::FpMontgomeryDomainFieldElement([0; 12]);
         fiat::fp_to_montgomery(&mut x, &x_non_monty);
@@ -38,30 +46,38 @@ impl Fp {
         Self(x)
     }
 
-    pub fn to_le_limbs(&self) -> [u32; 12] {
+    pub fn to_le_limbs(&self) -> [u64; 6] {
         let mut x_non_montgomery = fiat::FpNonMontgomeryDomainFieldElement([0; 12]);
         fiat::fp_from_montgomery(&mut x_non_montgomery, &self.0);
-        x_non_montgomery.0
+        let limbs = x_non_montgomery.0;
+        let mut out = [0u64; 6];
+        for i in 0..6 {
+            out[i] = (limbs[2 * i] as u64) | ((limbs[2 * i + 1] as u64) << 32);
+        }
+        out
     }
 
     pub fn to_bytes_le(&self) -> [u8; 48] {
         let mut bytes = [0u8; 48];
-        fiat::fp_to_bytes(&mut bytes, &self.to_le_limbs());
+        let mut x_non_montgomery = fiat::FpNonMontgomeryDomainFieldElement([0; 12]);
+        fiat::fp_from_montgomery(&mut x_non_montgomery, &self.0);
+        fiat::fp_to_bytes(&mut bytes, &x_non_montgomery.0);
         bytes
     }
 
     pub const fn from_montgomery_limbs(limbs: [u32; 12]) -> Fp {
         Self(fiat::FpMontgomeryDomainFieldElement(limbs))
     }
 
-    pub fn zero() -> Fp {
+    pub const fn zero() -> Fp {
         Self(fiat::FpMontgomeryDomainFieldElement([0; 12]))
     }
 
-    pub fn one() -> Fp {
-        let mut one = Self::zero();
-        fiat::fp_set_one(&mut one.0);
-        one
+    pub const fn one() -> Self {
+        Self(fiat::FpMontgomeryDomainFieldElement([
+            4294967144, 47054847, 2147483569, 1363189635, 2323464178, 2675815337, 1853645573,
+            2068748215, 2151449832, 1291097535, 3808294042, 9266785,
+        ]))
     }
 
     pub fn square(&self) -> Fp {
@@ -149,19 +165,19 @@ impl Fp {
         Some(Fp(result))
     }
 
-    pub fn add(self, other: Fp) -> Fp {
+    pub fn add(self, other: &Fp) -> Fp {
         let mut result = fiat::FpMontgomeryDomainFieldElement([0; 12]);
         fiat::fp_add(&mut result, &self.0, &other.0);
         Fp(result)
     }
 
-    pub fn sub(self, other: Fp) -> Fp {
+    pub fn sub(self, other: &Fp) -> Fp {
         let mut result = fiat::FpMontgomeryDomainFieldElement([0; 12]);
         fiat::fp_sub(&mut result, &self.0, &other.0);
         Fp(result)
     }
 
-    pub fn mul(self, other: Fp) -> Fp {
+    pub fn mul(self, other: &Fp) -> Fp {
         let mut result = fiat::FpMontgomeryDomainFieldElement([0; 12]);
         fiat::fp_mul(&mut result, &self.0, &other.0);
         Fp(result)