PackedField trait

src/field/mod.rs

@@ -4,6 +4,8 @@ pub mod extension_field;
 pub mod fft;
 pub mod field_types;
 pub(crate) mod interpolation;
+pub(crate) mod packable;
+pub(crate) mod packed_field;
 
 #[cfg(test)]
 mod field_testing;

src/field/packable.rs

@@ -0,0 +1,14 @@
+use crate::field::crandall_field::CrandallField;
+use crate::field::field_types::Field;
+use crate::field::packed_field::{PackedField, Singleton};
+
+/// Points us to the default packing for a particular field. There may me multiple choices of
+/// PackedField for a particular Field (e.g. Singleton works for all fields), but this is the
+/// recommended one. The recommended packing varies by target_arch and target_feature.
+pub trait Packable: Field {
+    type PackedType: PackedField<FieldType = Self>;
+}
+
+impl<F: Field> Packable for F {
+    default type PackedType = Singleton<Self>;
+}

src/field/packed_field.rs

@@ -0,0 +1,211 @@
+use std::fmt;
+use std::fmt::{Debug, Formatter};
+use std::iter::{Product, Sum};
+use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign};
+
+use crate::field::field_types::Field;
+
+pub trait PackedField:
+    'static
+    + Add<Self, Output = Self>
+    + Add<Self::FieldType, Output = Self>
+    + AddAssign<Self>
+    + AddAssign<Self::FieldType>
+    + Copy
+    + Debug
+    + Default
+    // TODO: Implementing Div sounds like a pain so it's a worry for later.
+    + Mul<Self, Output = Self>
+    + Mul<Self::FieldType, Output = Self>
+    + MulAssign<Self>
+    + MulAssign<Self::FieldType>
+    + Neg<Output = Self>
+    + Product
+    + Send
+    + Sub<Self, Output = Self>
+    + Sub<Self::FieldType, Output = Self>
+    + SubAssign<Self>
+    + SubAssign<Self::FieldType>
+    + Sum
+    + Sync
+{
+    type FieldType: Field;
+
+    const LOG2_WIDTH: usize;
+    const WIDTH: usize = 1 << Self::LOG2_WIDTH;
+
+    fn square(&self) -> Self;
+
+    fn zero() -> Self;
+    fn one() -> Self;
+
+    fn broadcast(x: Self::FieldType) -> Self;
+
+    fn new_from_slice(arr: &[Self::FieldType]) -> Self;
+    fn to_vec(&self) -> Vec<Self::FieldType>;
+
+    fn interleave(&self, other: Self, r: usize) -> (Self, Self);
+
+    fn pack_slice(buf: &[Self::FieldType]) -> &[Self] {
+        assert!(
+            buf.len() % Self::WIDTH == 0,
+            "Slice length (got {}) must be a multiple of packed field width ({}).",
+            buf.len(),
+            Self::WIDTH
+        );
+        let buf_ptr = buf.as_ptr().cast::<Self>();
+        // assert!(buf_ptr.align_offset(align_of::<Self>()) == 0, "Not aligned");
+        let n = buf.len() / Self::WIDTH;
+        unsafe { std::slice::from_raw_parts(buf_ptr, n) }
+    }
+    fn pack_slice_mut(buf: &mut [Self::FieldType]) -> &mut [Self] {
+        assert!(
+            buf.len() % Self::WIDTH == 0,
+            "Slice length (got {}) must be a multiple of packed field width ({}).",
+            buf.len(),
+            Self::WIDTH
+        );
+        let buf_ptr = buf.as_mut_ptr().cast::<Self>();
+        // assert!(buf_ptr.align_offset(align_of::<Self>()) == 0, "Not aligned");
+        let n = buf.len() / Self::WIDTH;
+        unsafe { std::slice::from_raw_parts_mut(buf_ptr, n) }
+    }
+}
+
+#[derive(Copy, Clone)]
+#[repr(transparent)]
+pub struct Singleton<F: Field>(pub F);
+
+impl<F: Field> Add<Self> for Singleton<F> {
+    type Output = Self;
+    fn add(self, rhs: Self) -> Self {
+        Self(self.0 + rhs.0)
+    }
+}
+impl<F: Field> Add<F> for Singleton<F> {
+    type Output = Self;
+    fn add(self, rhs: F) -> Self {
+        self + Self::broadcast(rhs)
+    }
+}
+impl<F: Field> AddAssign<Self> for Singleton<F> {
+    fn add_assign(&mut self, rhs: Self) {
+        *self = *self + rhs;
+    }
+}
+impl<F: Field> AddAssign<F> for Singleton<F> {
+    fn add_assign(&mut self, rhs: F) {
+        *self = *self + rhs;
+    }
+}
+
+impl<F: Field> Debug for Singleton<F> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "({:?})", self.0)
+    }
+}
+
+impl<F: Field> Default for Singleton<F> {
+    fn default() -> Self {
+        Self::zero()
+    }
+}
+
+impl<F: Field> Mul<Self> for Singleton<F> {
+    type Output = Self;
+    fn mul(self, rhs: Self) -> Self {
+        Self(self.0 * rhs.0)
+    }
+}
+impl<F: Field> Mul<F> for Singleton<F> {
+    type Output = Self;
+    fn mul(self, rhs: F) -> Self {
+        self * Self::broadcast(rhs)
+    }
+}
+impl<F: Field> MulAssign<Self> for Singleton<F> {
+    fn mul_assign(&mut self, rhs: Self) {
+        *self = *self * rhs;
+    }
+}
+impl<F: Field> MulAssign<F> for Singleton<F> {
+    fn mul_assign(&mut self, rhs: F) {
+        *self = *self * rhs;
+    }
+}
+
+impl<F: Field> Neg for Singleton<F> {
+    type Output = Self;
+    fn neg(self) -> Self {
+        Self(-self.0)
+    }
+}
+
+impl<F: Field> Product for Singleton<F> {
+    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
+        Self(iter.map(|x| x.0).product())
+    }
+}
+
+impl<F: Field> PackedField for Singleton<F> {
+    const LOG2_WIDTH: usize = 0;
+    type FieldType = F;
+
+    fn square(&self) -> Self {
+        *self * *self
+    }
+
+    fn zero() -> Self {
+        Self::broadcast(F::ZERO)
+    }
+    fn one() -> Self {
+        Self::broadcast(F::ONE)
+    }
+
+    fn broadcast(x: F) -> Self {
+        Self(x)
+    }
+
+    fn new_from_slice(arr: &[Self::FieldType]) -> Self {
+        Self(arr[0])
+    }
+    fn to_vec(&self) -> Vec<Self::FieldType> {
+        vec![self.0]
+    }
+
+    fn interleave(&self, other: Self, r: usize) -> (Self, Self) {
+        match r {
+            0 => (*self, other), // This is a no-op whenever r == LOG2_WIDTH.
+            _ => panic!("r cannot be more than LOG2_WIDTH"),
+        }
+    }
+}
+
+impl<F: Field> Sub<Self> for Singleton<F> {
+    type Output = Self;
+    fn sub(self, rhs: Self) -> Self {
+        Self(self.0 - rhs.0)
+    }
+}
+impl<F: Field> Sub<F> for Singleton<F> {
+    type Output = Self;
+    fn sub(self, rhs: F) -> Self {
+        self - Self::broadcast(rhs)
+    }
+}
+impl<F: Field> SubAssign<Self> for Singleton<F> {
+    fn sub_assign(&mut self, rhs: Self) {
+        *self = *self - rhs;
+    }
+}
+impl<F: Field> SubAssign<F> for Singleton<F> {
+    fn sub_assign(&mut self, rhs: F) {
+        *self = *self - rhs;
+    }
+}
+
+impl<F: Field> Sum for Singleton<F> {
+    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
+        Self(iter.map(|x| x.0).sum())
+    }
+}

src/lib.rs

@@ -1,5 +1,6 @@
 #![feature(destructuring_assignment)]
 #![feature(generic_const_exprs)]
+#![feature(specialization)]
 
 pub mod field;
 pub mod fri;