From aa73d7b1bc16147879ae37761c44e95b464a6ac4 Mon Sep 17 00:00:00 2001
From: Henry de Valence <hdevalence@hdevalence.ca>
Date: Sun, 11 Nov 2018 21:36:55 -0800
Subject: [PATCH] Add a FieldElement51x4 type.

---
 src/backend/vector/ifma/field.rs | 115 ++++++++++++++++++++++++++++++-
 src/backend/vector/ifma/mod.rs   |   2 +-
 src/lib.rs                       |   1 -
 3 files changed, 113 insertions(+), 5 deletions(-)

diff --git a/src/backend/vector/ifma/field.rs b/src/backend/vector/ifma/field.rs
index 6ff36b7..d2f59e8 100644
--- a/src/backend/vector/ifma/field.rs
+++ b/src/backend/vector/ifma/field.rs
@@ -20,6 +20,83 @@ extern "C" {
     fn madd52hi(z: u64x4, x: u64x4, y: u64x4) -> u64x4;
 }
 
+pub struct FieldElement51x4([u64x4; 5]);
+
+impl FieldElement51x4 {
+    pub fn new(
+        x0: &FieldElement51,
+        x1: &FieldElement51,
+        x2: &FieldElement51,
+        x3: &FieldElement51,
+    ) -> FieldElement51x4 {
+        FieldElement51x4([
+            u64x4::new(x0.0[0], x1.0[0], x2.0[0], x3.0[0]),
+            u64x4::new(x0.0[1], x1.0[1], x2.0[1], x3.0[1]),
+            u64x4::new(x0.0[2], x1.0[2], x2.0[2], x3.0[2]),
+            u64x4::new(x0.0[3], x1.0[3], x2.0[3], x3.0[3]),
+            u64x4::new(x0.0[4], x1.0[4], x2.0[4], x3.0[4]),
+        ])
+        .reduce()
+    }
+
+    pub fn split(&self) -> [FieldElement51; 4] {
+        let x = &self.0;
+        [
+            FieldElement51([
+                x[0].extract(0),
+                x[1].extract(0),
+                x[2].extract(0),
+                x[3].extract(0),
+                x[4].extract(0),
+            ]),
+            FieldElement51([
+                x[0].extract(1),
+                x[1].extract(1),
+                x[2].extract(1),
+                x[3].extract(1),
+                x[4].extract(1),
+            ]),
+            FieldElement51([
+                x[0].extract(2),
+                x[1].extract(2),
+                x[2].extract(2),
+                x[3].extract(2),
+                x[4].extract(2),
+            ]),
+            FieldElement51([
+                x[0].extract(3),
+                x[1].extract(3),
+                x[2].extract(3),
+                x[3].extract(3),
+                x[4].extract(3),
+            ]),
+        ]
+    }
+
+    #[inline]
+    pub fn reduce(&self) -> FieldElement51x4 {
+        let mask = u64x4::splat((1 << 51) - 1);
+        let r19 = u64x4::splat(19);
+
+        // Compute carryouts in parallel
+        let c0 = self.0[0] >> 51;
+        let c1 = self.0[1] >> 51;
+        let c2 = self.0[2] >> 51;
+        let c3 = self.0[3] >> 51;
+        let c4 = self.0[4] >> 51;
+
+        unsafe {
+            FieldElement51x4([
+                madd52lo(self.0[0] & mask, c4, r19),
+                (self.0[1] & mask) + c0,
+                (self.0[2] & mask) + c1,
+                (self.0[3] & mask) + c2,
+                (self.0[4] & mask) + c3,
+            ])
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -32,9 +109,41 @@ mod test {
 
         z = unsafe { madd52lo(z, x, y) };
 
-        assert_eq!(z, u64x4::splat(5 + 2*3));
+        assert_eq!(z, u64x4::splat(5 + 2 * 3));
+    }
+
+    #[test]
+    fn new_split_round_trip_on_reduced_input() {
+        // Invert a small field element to get a big one
+        let a = FieldElement51([2438, 24, 243, 0, 0]).invert();
+
+        let ax4 = FieldElement51x4::new(&a, &a, &a, &a);
+        let splits = ax4.split();
+
+        for i in 0..4 {
+            assert_eq!(a, splits[i]);
+        }
+    }
+
+    #[test]
+    fn new_split_round_trip_on_unreduced_input() {
+        // Invert a small field element to get a big one
+        let a = FieldElement51([2438, 24, 243, 0, 0]).invert();
+        // ... but now multiply it by 16 without reducing coeffs
+        let a16 = FieldElement51([
+            a.0[0] << 4,
+            a.0[1] << 4,
+            a.0[2] << 4,
+            a.0[3] << 4,
+            a.0[4] << 4,
+        ]);
+
+        let a16x4 = FieldElement51x4::new(&a16, &a16, &a16, &a16);
+        let splits = a16x4.split();
+
+        for i in 0..4 {
+            assert_eq!(a16, splits[i]);
+        }
     }
 
 }
-
-
diff --git a/src/backend/vector/ifma/mod.rs b/src/backend/vector/ifma/mod.rs
index 4238412..e066ddd 100644
--- a/src/backend/vector/ifma/mod.rs
+++ b/src/backend/vector/ifma/mod.rs
@@ -7,5 +7,5 @@
 // Authors:
 // - Henry de Valence <hdevalence@hdevalence.ca>
 
-pub(crate) mod field;
+pub mod field;
 
diff --git a/src/lib.rs b/src/lib.rs
index 2bf5c8a..9dd6608 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -101,4 +101,3 @@ pub(crate) mod prelude;
 
 // Generic code for window lookups
 pub(crate) mod window;
-