diff --git a/ext/crates/fp/src/matrix/matrix_inner.rs b/ext/crates/fp/src/matrix/matrix_inner.rs index a2b82c0a0..c8ef3a866 100644 --- a/ext/crates/fp/src/matrix/matrix_inner.rs +++ b/ext/crates/fp/src/matrix/matrix_inner.rs @@ -367,12 +367,12 @@ impl Matrix { pub fn row(&self, row: usize) -> FpSlice<'_> { let limb_range = row_to_limb_range(row, self.stride); - FpSlice::new(self.prime(), &self.data[limb_range], 0, self.columns) + FpSlice::_new(self.prime(), &self.data[limb_range], 0, self.columns) } pub fn row_mut(&mut self, row: usize) -> FpSliceMut<'_> { let limb_range = row_to_limb_range(row, self.stride); - FpSliceMut::new(self.prime(), &mut self.data[limb_range], 0, self.columns) + FpSliceMut::_new(self.prime(), &mut self.data[limb_range], 0, self.columns) } } @@ -393,7 +393,7 @@ impl Matrix { .data .chunks_mut(self.stride) .take(logical_rows) // Only iterate over logical rows - .map(move |row| FpSliceMut::new(p, row, 0, columns)); + .map(move |row| FpSliceMut::_new(p, row, 0, columns)); Either::Right(rows) } } @@ -412,7 +412,7 @@ impl Matrix { .data .maybe_par_chunks_mut(self.stride) .take(logical_rows) // Only iterate over logical rows - .map(move |row| FpSliceMut::new(p, row, 0, columns)); + .map(move |row| FpSliceMut::_new(p, row, 0, columns)); Either::Right(rows) } } @@ -516,8 +516,8 @@ impl Matrix { let row1 = unsafe { std::slice::from_raw_parts_mut(ptr.add(i * self.stride), self.stride) }; let row2 = unsafe { std::slice::from_raw_parts_mut(ptr.add(j * self.stride), self.stride) }; ( - FpSliceMut::new(self.prime(), row1, 0, self.columns), - FpSliceMut::new(self.prime(), row2, 0, self.columns), + FpSliceMut::_new(self.prime(), row1, 0, self.columns), + FpSliceMut::_new(self.prime(), row2, 0, self.columns), ) } @@ -1263,7 +1263,7 @@ impl AugmentedMatrix { let start_idx = self.start[start]; let end_idx = self.end[end]; let limb_range = row_to_limb_range(i, self.stride); - FpSliceMut::new(self.prime(), &mut self.data[limb_range], start_idx, end_idx) + FpSliceMut::_new(self.prime(), &mut self.data[limb_range], start_idx, end_idx) } pub fn row_segment(&self, i: usize, start: usize, end: usize) -> FpSlice<'_> { @@ -1414,7 +1414,7 @@ impl<'a> MatrixSliceMut<'a> { let end = self.col_end; (0..self.rows).map(move |row_idx| { let limb_range = row_to_limb_range(row_idx, self.stride); - FpSlice::new(self.prime(), &self.data[limb_range], start, end) + FpSlice::_new(self.prime(), &self.data[limb_range], start, end) }) } @@ -1429,7 +1429,7 @@ impl<'a> MatrixSliceMut<'a> { let rows = self .data .chunks_mut(self.stride) - .map(move |row| FpSliceMut::new(p, row, start, end)); + .map(move |row| FpSliceMut::_new(p, row, start, end)); Either::Right(rows) } } @@ -1447,14 +1447,14 @@ impl<'a> MatrixSliceMut<'a> { let rows = self .data .maybe_par_chunks_mut(self.stride) - .map(move |row| FpSliceMut::new(p, row, start, end)); + .map(move |row| FpSliceMut::_new(p, row, start, end)); Either::Right(rows) } } pub fn row(&mut self, row: usize) -> FpSlice<'_> { let limb_range = row_to_limb_range(row, self.stride); - FpSlice::new( + FpSlice::_new( self.prime(), &self.data[limb_range], self.col_start, @@ -1464,7 +1464,7 @@ impl<'a> MatrixSliceMut<'a> { pub fn row_mut(&mut self, row: usize) -> FpSliceMut<'_> { let limb_range = row_to_limb_range(row, self.stride); - FpSliceMut::new( + FpSliceMut::_new( self.prime(), &mut self.data[limb_range], self.col_start, diff --git a/ext/crates/fp/src/vector/fp_wrapper/helpers.rs b/ext/crates/fp/src/vector/fp_wrapper/helpers.rs index e6e7f7e17..fa3068548 100644 --- a/ext/crates/fp/src/vector/fp_wrapper/helpers.rs +++ b/ext/crates/fp/src/vector/fp_wrapper/helpers.rs @@ -14,38 +14,36 @@ use itertools::Itertools; -use super::{FqSlice, FqSliceMut, FqVector, FqVectorIterator, FqVectorNonZeroIterator}; +use super::{ + FqSlice, FqSliceMut, FqVector, FqVectorBase, FqVectorIterator, FqVectorNonZeroIterator, Repr, + ReprMut, +}; use crate::field::Field; -impl FqVector { - pub(super) fn scale_helper(&mut self, c: F::ElementContainer) { - self.scale(self.fq().el(c)) - } - +impl FqVectorBase { pub(super) fn entry_helper(&self, index: usize) -> F::ElementContainer { self.entry(index).val() } +} - pub(super) fn set_entry_helper(&mut self, index: usize, value: F::ElementContainer) { - self.set_entry(index, self.fq().el(value)) - } - - pub(super) fn add_helper(&mut self, other: &Self, c: F::ElementContainer) { - self.add(other, self.fq().el(c)) +impl FqVectorBase { + pub(super) fn scale_helper(&mut self, c: F::ElementContainer) { + self.scale(self.fq().el(c)) } - pub(super) fn add_offset_helper( - &mut self, - other: &Self, - c: F::ElementContainer, - offset: usize, - ) { - self.add_offset(other, self.fq().el(c), offset) + pub(super) fn set_entry_helper(&mut self, index: usize, value: F::ElementContainer) { + self.set_entry(index, self.fq().el(value)) } pub(super) fn add_basis_element_helper(&mut self, index: usize, value: F::ElementContainer) { self.add_basis_element(index, self.fq().el(value)) } +} + +impl FqVector { + pub(super) fn add_helper(&mut self, other: &Self, c: F::ElementContainer) { + self.add(other, self.fq().el(c)) + } pub(super) fn copy_from_slice_helper(&mut self, other: &[F::ElementContainer]) { self.copy_from_slice(&other.iter().map(|x| self.fq().el(x.clone())).collect_vec()) @@ -71,26 +69,27 @@ impl FqVector { self.add_carry(other, self.fq().el(c), rest) } + pub(super) fn add_offset_helper( + &mut self, + other: &Self, + c: F::ElementContainer, + offset: usize, + ) { + self.add_offset(other, self.fq().el(c), offset) + } + pub(super) fn first_nonzero_helper(&self) -> Option<(usize, F::ElementContainer)> { self.first_nonzero().map(|(idx, c)| (idx, c.val())) } } impl FqSlice<'_, F> { - pub(super) fn entry_helper(&self, index: usize) -> F::ElementContainer { - self.entry(index).val() - } - pub(super) fn first_nonzero_helper(&self) -> Option<(usize, F::ElementContainer)> { self.first_nonzero().map(|(idx, c)| (idx, c.val())) } } impl FqSliceMut<'_, F> { - pub(super) fn scale_helper(&mut self, c: F::ElementContainer) { - self.scale(self.fq().el(c)) - } - pub(super) fn add_helper(&mut self, other: FqSlice, c: F::ElementContainer) { self.add(other, self.fq().el(c)) } @@ -104,14 +103,6 @@ impl FqSliceMut<'_, F> { self.add_offset(other, self.fq().el(c), offset) } - pub(super) fn set_entry_helper(&mut self, index: usize, value: F::ElementContainer) { - self.set_entry(index, self.fq().el(value)) - } - - pub(super) fn add_basis_element_helper(&mut self, index: usize, value: F::ElementContainer) { - self.add_basis_element(index, self.fq().el(value)) - } - pub(super) fn add_masked_helper( &mut self, other: FqSlice, diff --git a/ext/crates/fp/src/vector/fp_wrapper/mod.rs b/ext/crates/fp/src/vector/fp_wrapper/mod.rs index 9ba2eca5c..7238dacbb 100644 --- a/ext/crates/fp/src/vector/fp_wrapper/mod.rs +++ b/ext/crates/fp/src/vector/fp_wrapper/mod.rs @@ -17,12 +17,14 @@ use std::{convert::TryInto, io}; use itertools::Itertools; use serde::{Deserialize, Deserializer, Serialize, Serializer}; -use super::iter::{FqVectorIterator, FqVectorNonZeroIterator}; +use super::{ + FqSlice, FqSliceMut, FqVector, FqVectorBase, Repr, ReprMut, + iter::{FqVectorIterator, FqVectorNonZeroIterator}, +}; use crate::{ field::{Fp, field_internal::FieldInternal}, limb::Limb, prime::Prime, - vector::inner::{FqSlice, FqSliceMut, FqVector}, }; mod helpers; @@ -43,20 +45,22 @@ use macros_2::{dispatch_struct, dispatch_vector, impl_try_into, use_primes}; use_primes!(); -dispatch_struct! { - #[derive(Debug, Hash, Eq, PartialEq, Clone)] - pub FpVector from FqVector -} - -dispatch_struct! { - #[derive(Debug, Copy, Clone)] - pub FpSlice<'a> from FqSlice +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] +pub enum FpVectorBase { + _2(FqVectorBase>), + #[cfg(feature = "odd-primes")] + _3(FqVectorBase>), + #[cfg(feature = "odd-primes")] + _5(FqVectorBase>), + #[cfg(feature = "odd-primes")] + _7(FqVectorBase>), + #[cfg(feature = "odd-primes")] + Big(FqVectorBase>), } -dispatch_struct! { - #[derive(Debug)] - pub FpSliceMut<'a> from FqSliceMut -} +pub type FpVector = FpVectorBase>; +pub type FpSlice<'a> = FpVectorBase; +pub type FpSliceMut<'a> = FpVectorBase; dispatch_struct! { pub FpVectorIterator<'a> from FqVectorIterator @@ -66,29 +70,43 @@ dispatch_struct! { pub FpVectorNonZeroIterator<'a> from FqVectorNonZeroIterator } -impl FpVector { +impl FpVectorBase { dispatch_vector! { pub fn prime(&self) -> ValidPrime; pub fn len(&self) -> usize; pub fn is_empty(&self) -> bool; - pub fn @scale(&mut self, c: u32); - pub fn set_to_zero(&mut self); + pub fn slice(&self, start: usize, end: usize) -> (dispatch FpSlice<'_>); + pub fn as_slice(&self) -> (dispatch FpSlice<'_>); + pub fn iter(&self) -> (dispatch FpVectorIterator<'_>); + pub fn iter_nonzero(&self) -> (dispatch FpVectorNonZeroIterator<'_>); + pub fn is_zero(&self) -> bool; pub fn @entry(&self, index: usize) -> u32; + + pub(crate) fn limbs(&self) -> (&[Limb]); + } +} + +impl FpVectorBase { + dispatch_vector! { + pub fn slice_mut(&mut self, start: usize, end: usize) -> (dispatch FpSliceMut<'_>); + pub fn as_slice_mut(&mut self) -> (dispatch FpSliceMut<'_>); + pub fn set_to_zero(&mut self); + pub fn @scale(&mut self, c: u32); pub fn @set_entry(&mut self, index: usize, value: u32); + pub fn @add_basis_element(&mut self, index: usize, value: u32); + + pub(crate) fn limbs_mut(&mut self) -> (&mut [Limb]); + } +} + +impl FpVector { + dispatch_vector! { pub fn assign(&mut self, other: &Self); pub fn assign_partial(&mut self, other: &Self); pub fn @add(&mut self, other: &Self, c: u32); pub fn @add_offset(&mut self, other: &Self, c: u32, offset: usize); - pub fn slice(&self, start: usize, end: usize) -> (dispatch FpSlice<'_>); - pub fn as_slice(&self) -> (dispatch FpSlice<'_>); - pub fn slice_mut(&mut self, start: usize, end: usize) -> (dispatch FpSliceMut<'_>); - pub fn as_slice_mut(&mut self) -> (dispatch FpSliceMut<'_>); - pub fn is_zero(&self) -> bool; - pub fn iter(&self) -> (dispatch FpVectorIterator<'_>); - pub fn iter_nonzero(&self) -> (dispatch FpVectorNonZeroIterator<'_>); pub fn extend_len(&mut self, dim: usize); pub fn set_scratch_vector_size(&mut self, dim: usize); - pub fn @add_basis_element(&mut self, index: usize, value: u32); pub fn @copy_from_slice(&mut self, slice: &[u32]); pub fn @add_truncate(&mut self, other: &Self, c: u32) -> (Option<()>); pub fn sign_rule(&self, other: &Self) -> bool; @@ -96,8 +114,6 @@ impl FpVector { pub fn @first_nonzero(&self) -> (Option<(usize, u32)>); pub fn density(&self) -> f32; - pub(crate) fn limbs(&self) -> (&[Limb]); - pub fn new(p: P, len: usize) -> (from FqVector); pub fn new_with_capacity(p: P, len: usize, capacity: usize) -> (from FqVector); @@ -125,43 +141,24 @@ impl FpVector { impl<'a> FpSlice<'a> { dispatch_vector! { - pub(crate) fn new(p: P, limbs: &'a [Limb], start: usize, end: usize) -> (from FqSlice); - pub fn prime(&self) -> ValidPrime; - pub fn len(&self) -> usize; - pub fn is_empty(&self) -> bool; - pub fn @entry(&self, index: usize) -> u32; - pub fn iter(self) -> (dispatch FpVectorIterator<'a>); - pub fn iter_nonzero(self) -> (dispatch FpVectorNonZeroIterator<'a>); + pub(crate) fn _new(p: P, limbs: &'a [Limb], start: usize, end: usize) -> (from FqSlice); pub fn @first_nonzero(&self) -> (Option<(usize, u32)>); - pub fn is_zero(&self) -> bool; pub fn restrict(self, start: usize, end: usize) -> (dispatch FpSlice<'a>); pub fn to_owned(self) -> (dispatch FpVector); - - pub(crate) fn limbs(&self) -> (&[Limb]); } } impl<'a> FpSliceMut<'a> { dispatch_vector! { - pub(crate) fn new(p: P, limbs: &'a mut [Limb], start: usize, end: usize) -> (from FqSliceMut); - pub fn prime(&self) -> ValidPrime; - pub fn @scale(&mut self, c: u32); - pub fn set_to_zero(&mut self); + pub(crate) fn _new(p: P, limbs: &'a mut [Limb], start: usize, end: usize) -> (from FqSliceMut); pub fn @add(&mut self, other: FpSlice, c: u32); pub fn @add_offset(&mut self, other: FpSlice, c: u32, offset: usize); pub fn assign(&mut self, other: FpSlice); pub fn shl_assign(&mut self, shift: usize); - pub fn @set_entry(&mut self, index: usize, value: u32); - pub fn as_slice(&self) -> (dispatch FpSlice<'_>); - pub fn slice_mut(&mut self, start: usize, end: usize) -> (dispatch FpSliceMut<'_>); - pub fn @add_basis_element(&mut self, index: usize, value: u32); pub fn copy(&mut self) -> (dispatch FpSliceMut<'_>); pub fn @add_masked(&mut self, other: FpSlice, c: u32, mask: &[usize]); pub fn @add_unmasked(&mut self, other: FpSlice, c: u32, mask: &[usize]); pub fn @add_tensor(&mut self, offset: usize, coeff: u32, @left: FpSlice, right: FpSlice); - - pub(crate) fn limbs(&self) -> (&[Limb]); - pub(crate) fn limbs_mut(&mut self) -> (&mut [Limb]); } } @@ -171,13 +168,7 @@ impl FpVectorIterator<'_> { } } -impl std::fmt::Display for FpVector { - fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - self.as_slice().fmt(f) - } -} - -impl std::fmt::Display for FpSlice<'_> { +impl std::fmt::Display for FpVectorBase { /// # Example /// ``` /// # use fp::vector::FpVector; @@ -199,8 +190,8 @@ impl std::fmt::Display for FpSlice<'_> { } } -impl From<&FpVector> for Vec { - fn from(v: &FpVector) -> Self { +impl From<&FpVectorBase> for Vec { + fn from(v: &FpVectorBase) -> Self { v.iter().collect() } } @@ -258,33 +249,15 @@ impl<'de> Deserialize<'de> for FpVector { } } -impl<'a, 'b> From<&'a mut FpSliceMut<'b>> for FpSliceMut<'a> { - fn from(slice: &'a mut FpSliceMut<'b>) -> Self { - slice.copy() - } -} - -impl<'a, 'b> From<&'a FpSlice<'b>> for FpSlice<'a> { - fn from(slice: &'a FpSlice<'b>) -> Self { - *slice - } -} - -impl<'a, 'b> From<&'a FpSliceMut<'b>> for FpSlice<'a> { - fn from(slice: &'a FpSliceMut<'b>) -> Self { - slice.as_slice() - } -} - -impl<'a> From<&'a FpVector> for FpSlice<'a> { - fn from(v: &'a FpVector) -> Self { - v.as_slice() +impl<'a, const A: bool, R: Repr> From<&'a FpVectorBase> for FpSlice<'a> { + fn from(value: &'a FpVectorBase) -> Self { + value.as_slice() } } -impl<'a> From<&'a mut FpVector> for FpSliceMut<'a> { - fn from(v: &'a mut FpVector) -> Self { - v.as_slice_mut() +impl<'a, const A: bool, R: ReprMut> From<&'a mut FpVectorBase> for FpSliceMut<'a> { + fn from(value: &'a mut FpVectorBase) -> Self { + value.as_slice_mut() } } diff --git a/ext/crates/fp/src/vector/impl_fqslice.rs b/ext/crates/fp/src/vector/impl_fqslice.rs index e2f243f0d..7dc90a77f 100644 --- a/ext/crates/fp/src/vector/impl_fqslice.rs +++ b/ext/crates/fp/src/vector/impl_fqslice.rs @@ -1,84 +1,18 @@ -use itertools::Itertools; - -use super::{ - inner::{FqSlice, FqVector}, - iter::{FqVectorIterator, FqVectorNonZeroIterator}, -}; -use crate::{ - constants, - field::{Field, element::FieldElement}, - limb::Limb, - prime::{Prime, ValidPrime}, -}; +use super::{FqSlice, FqVector, FqVectorBase, Repr}; +use crate::field::{Field, element::FieldElement}; // Public methods impl<'a, F: Field> FqSlice<'a, F> { - pub fn prime(&self) -> ValidPrime { - self.fq().characteristic().to_dyn() - } - - pub fn len(&self) -> usize { - self.end() - self.start() - } - - pub const fn is_empty(&self) -> bool { - self.start() == self.end() - } - - pub fn entry(&self, index: usize) -> FieldElement { - debug_assert!( - index < self.len(), - "Index {} too large, length of vector is only {}.", - index, - self.len() - ); - let bit_mask = self.fq().bitmask(); - let limb_index = self.fq().limb_bit_index_pair(index + self.start()); - let mut result = self.limbs()[limb_index.limb]; - result >>= limb_index.bit_index; - result &= bit_mask; - self.fq().decode(result) - } - - /// TODO: implement prime 2 version - pub fn iter(self) -> FqVectorIterator<'a, F> { - FqVectorIterator::new(self) - } - - pub fn iter_nonzero(self) -> FqVectorNonZeroIterator<'a, F> { - FqVectorNonZeroIterator::new(self) - } - pub fn first_nonzero(&self) -> Option<(usize, FieldElement)> { self.iter_nonzero().next() } - pub fn is_zero(&self) -> bool { - let limb_range = self.limb_range(); - if limb_range.is_empty() { - return true; - } - let (min_mask, max_mask) = self.limb_masks(); - if self.limbs()[limb_range.start] & min_mask != 0 { - return false; - } - - let inner_range = self.limb_range_inner(); - if !inner_range.is_empty() && self.limbs()[inner_range].iter().any(|&x| x != 0) { - return false; - } - if self.limbs()[limb_range.end - 1] & max_mask != 0 { - return false; - } - true - } - #[must_use] pub fn restrict(self, start: usize, end: usize) -> Self { assert!(start <= end && end <= self.len()); - FqSlice::new( + FqSlice::_new( self.fq(), self.into_limbs(), self.start() + start, @@ -104,84 +38,8 @@ impl<'a, F: Field> FqSlice<'a, F> { } } -// Limb methods -impl FqSlice<'_, F> { - #[inline] - pub(super) fn offset(&self) -> usize { - let bit_length = self.fq().bit_length(); - let entries_per_limb = self.fq().entries_per_limb(); - (self.start() % entries_per_limb) * bit_length - } - - #[inline] - pub(super) fn limb_range(&self) -> std::ops::Range { - self.fq().range(self.start(), self.end()) - } - - /// This function underflows if `self.end() == 0`, which happens if and only if we are taking a - /// slice of width 0 at the start of an `FpVector`. This should be a very rare edge case. - /// Dealing with the underflow properly would probably require using `saturating_sub` or - /// something of that nature, and that has a nontrivial (10%) performance hit. - #[inline] - pub(super) fn limb_range_inner(&self) -> std::ops::Range { - let range = self.limb_range(); - (range.start + 1)..(usize::max(range.start + 1, range.end - 1)) - } - - #[inline(always)] - pub(super) fn min_limb_mask(&self) -> Limb { - !0 << self.offset() - } - - #[inline(always)] - pub(super) fn max_limb_mask(&self) -> Limb { - let num_entries = 1 + (self.end() - 1) % self.fq().entries_per_limb(); - let bit_max = num_entries * self.fq().bit_length(); - - (!0) >> (constants::BITS_PER_LIMB - bit_max) - } - - #[inline(always)] - pub(super) fn limb_masks(&self) -> (Limb, Limb) { - if self.limb_range().len() == 1 { - ( - self.min_limb_mask() & self.max_limb_mask(), - self.min_limb_mask() & self.max_limb_mask(), - ) - } else { - (self.min_limb_mask(), self.max_limb_mask()) - } - } -} - -impl<'a, F: Field> From<&'a FqVector> for FqSlice<'a, F> { - fn from(v: &'a FqVector) -> Self { - v.slice(0, v.len()) - } -} - -impl std::fmt::Display for FqSlice<'_, F> { - /// # Example - /// ``` - /// # use fp::field::{Field, SmallFq}; - /// # use fp::prime::{P2, ValidPrime}; - /// # use fp::vector::FqVector; - /// let fq = SmallFq::new(P2, 3); - /// let v = FqVector::from_slice(fq, &[fq.zero(), fq.one(), fq.a(), fq.a() * fq.a()]); - /// assert_eq!(&format!("{v}"), "[0, 1, a, a^2]"); - /// - /// // This only looks reasonable over prime fields of order less than 10 - /// assert_eq!(&format!("{v:#}"), "01aa^2"); - /// ``` - fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - if f.alternate() { - for v in self.iter() { - // If self.p >= 11, this will look funky - write!(f, "{v}")?; - } - Ok(()) - } else { - write!(f, "[{}]", self.iter().format(", ")) - } +impl<'a, const A: bool, R: Repr, F: Field> From<&'a FqVectorBase> for FqSlice<'a, F> { + fn from(v: &'a FqVectorBase) -> Self { + v.as_slice() } } diff --git a/ext/crates/fp/src/vector/impl_fqslicemut.rs b/ext/crates/fp/src/vector/impl_fqslicemut.rs index d65d30247..4903e5cf9 100644 --- a/ext/crates/fp/src/vector/impl_fqslicemut.rs +++ b/ext/crates/fp/src/vector/impl_fqslicemut.rs @@ -2,121 +2,32 @@ use std::cmp::Ordering; use itertools::Itertools; -use super::inner::{FqSlice, FqSliceMut, FqVector}; use crate::{ constants, field::{Field, element::FieldElement}, limb::Limb, - prime::{Prime, ValidPrime}, + vector::{FqSlice, FqSliceMut, FqVectorBase, ReprMut}, }; -impl<'a, F: Field> FqSliceMut<'a, F> { - pub fn prime(&self) -> ValidPrime { - self.fq().characteristic().to_dyn() - } - - pub fn add_basis_element(&mut self, index: usize, value: FieldElement) { - assert_eq!(self.fq(), value.field()); - if self.fq().q() == 2 { - let pair = self.fq().limb_bit_index_pair(index + self.start()); - self.limbs_mut()[pair.limb] ^= self.fq().encode(value) << pair.bit_index; - } else { - let mut x = self.as_slice().entry(index); - x += value; - self.set_entry(index, x); - } - } - - pub fn set_entry(&mut self, index: usize, value: FieldElement) { - assert_eq!(self.fq(), value.field()); - assert!(index < self.as_slice().len()); - let bit_mask = self.fq().bitmask(); - let limb_index = self.fq().limb_bit_index_pair(index + self.start()); - let mut result = self.limbs()[limb_index.limb]; - result &= !(bit_mask << limb_index.bit_index); - result |= self.fq().encode(value) << limb_index.bit_index; - self.limbs_mut()[limb_index.limb] = result; - } - - fn reduce_limbs(&mut self) { - let fq = self.fq(); - if fq.q() != 2 { - let limb_range = self.as_slice().limb_range(); - - for limb in self.limbs_mut()[limb_range].iter_mut() { - *limb = fq.reduce(*limb); - } - } - } - - pub fn scale(&mut self, c: FieldElement) { - assert_eq!(self.fq(), c.field()); - let fq = self.fq(); - - if fq.q() == 2 { - if c == fq.zero() { - self.set_to_zero(); - } - return; - } - - let limb_range = self.as_slice().limb_range(); - if limb_range.is_empty() { - return; - } - let (min_mask, max_mask) = self.as_slice().limb_masks(); - - let limb = self.limbs()[limb_range.start]; - let masked_limb = limb & min_mask; - let rest_limb = limb & !min_mask; - self.limbs_mut()[limb_range.start] = fq.fma_limb(0, masked_limb, c.clone()) | rest_limb; - - let inner_range = self.as_slice().limb_range_inner(); - for limb in self.limbs_mut()[inner_range].iter_mut() { - *limb = fq.fma_limb(0, *limb, c.clone()); - } - if limb_range.len() > 1 { - let full_limb = self.limbs()[limb_range.end - 1]; - let masked_limb = full_limb & max_mask; - let rest_limb = full_limb & !max_mask; - self.limbs_mut()[limb_range.end - 1] = fq.fma_limb(0, masked_limb, c) | rest_limb; - } - self.reduce_limbs(); - } - - pub fn set_to_zero(&mut self) { - let limb_range = self.as_slice().limb_range(); - if limb_range.is_empty() { - return; - } - let (min_mask, max_mask) = self.as_slice().limb_masks(); - self.limbs_mut()[limb_range.start] &= !min_mask; - - let inner_range = self.as_slice().limb_range_inner(); - for limb in self.limbs_mut()[inner_range].iter_mut() { - *limb = 0; - } - self.limbs_mut()[limb_range.end - 1] &= !max_mask; - } - +impl FqSliceMut<'_, F> { pub fn add(&mut self, other: FqSlice<'_, F>, c: FieldElement) { assert_eq!(self.fq(), c.field()); assert_eq!(self.fq(), other.fq()); - if self.as_slice().is_empty() { + if self.is_empty() { return; } if self.fq().q() == 2 { if c != self.fq().zero() { - match self.as_slice().offset().cmp(&other.offset()) { + match self.offset().cmp(&other.offset()) { Ordering::Equal => self.add_shift_none(other, self.fq().one()), Ordering::Less => self.add_shift_left(other, self.fq().one()), Ordering::Greater => self.add_shift_right(other, self.fq().one()), }; } } else { - match self.as_slice().offset().cmp(&other.offset()) { + match self.offset().cmp(&other.offset()) { Ordering::Equal => self.add_shift_none(other, c), Ordering::Less => self.add_shift_left(other, c), Ordering::Greater => self.add_shift_right(other, c), @@ -153,12 +64,12 @@ impl<'a, F: Field> FqSliceMut<'a, F> { /// TODO: improve efficiency pub fn assign(&mut self, other: FqSlice<'_, F>) { assert_eq!(self.fq(), other.fq()); - if self.as_slice().offset() != other.offset() { + if self.offset() != other.offset() { self.set_to_zero(); self.add(other, self.fq().one()); return; } - let target_range = self.as_slice().limb_range(); + let target_range = self.limb_range(); let source_range = other.limb_range(); if target_range.is_empty() { @@ -171,7 +82,7 @@ impl<'a, F: Field> FqSliceMut<'a, F> { self.limbs_mut()[target_range.start] &= !min_mask; self.limbs_mut()[target_range.start] |= result; - let target_inner_range = self.as_slice().limb_range_inner(); + let target_inner_range = self.limb_range_inner(); let source_inner_range = other.limb_range_inner(); if !target_inner_range.is_empty() && !source_inner_range.is_empty() { self.limbs_mut()[target_inner_range] @@ -206,7 +117,7 @@ impl<'a, F: Field> FqSliceMut<'a, F> { assert_eq!(self.fq(), other.fq()); let fq = self.fq(); - let target_range = self.as_slice().limb_range(); + let target_range = self.limb_range(); let source_range = other.limb_range(); let (min_mask, max_mask) = other.limb_masks(); @@ -218,7 +129,7 @@ impl<'a, F: Field> FqSliceMut<'a, F> { ); self.limbs_mut()[target_range.start] = fq.reduce(self.limbs()[target_range.start]); - let target_inner_range = self.as_slice().limb_range_inner(); + let target_inner_range = self.limb_range_inner(); let source_inner_range = other.limb_range_inner(); if !source_inner_range.is_empty() { for (left, right) in self.limbs_mut()[target_inner_range] @@ -498,7 +409,7 @@ impl<'a, F: Field> FqSliceMut<'a, F> { // TODO: If this ends up being a bottleneck, try to use PDEP/PEXT assert_eq!(self.fq(), c.field()); assert_eq!(self.fq(), other.fq()); - assert_eq!(self.as_slice().len(), mask.len()); + assert_eq!(self.len(), mask.len()); for (i, &x) in mask.iter().enumerate() { let entry = other.entry(x); if entry != self.fq().zero() { @@ -517,24 +428,6 @@ impl<'a, F: Field> FqSliceMut<'a, F> { } } - pub fn slice_mut(&mut self, start: usize, end: usize) -> FqSliceMut<'_, F> { - assert!(start <= end && end <= self.as_slice().len()); - let orig_start = self.start(); - - FqSliceMut::new( - self.fq(), - self.limbs_mut(), - orig_start + start, - orig_start + end, - ) - } - - #[inline] - #[must_use] - pub fn as_slice(&self) -> FqSlice<'_, F> { - FqSlice::new(self.fq(), self.limbs(), self.start(), self.end()) - } - /// Generates a version of itself with a shorter lifetime #[inline] #[must_use] @@ -542,12 +435,14 @@ impl<'a, F: Field> FqSliceMut<'a, F> { let start = self.start(); let end = self.end(); - FqSliceMut::new(self.fq(), self.limbs_mut(), start, end) + FqSliceMut::_new(self.fq(), self.limbs_mut(), start, end) } } -impl<'a, F: Field> From<&'a mut FqVector> for FqSliceMut<'a, F> { - fn from(v: &'a mut FqVector) -> Self { - v.slice_mut(0, v.len()) +impl<'a, const A: bool, R: ReprMut, F: Field> From<&'a mut FqVectorBase> + for FqSliceMut<'a, F> +{ + fn from(v: &'a mut FqVectorBase) -> Self { + v.as_slice_mut() } } diff --git a/ext/crates/fp/src/vector/impl_fqvector.rs b/ext/crates/fp/src/vector/impl_fqvector.rs index 2a472a8b5..6a85147f7 100644 --- a/ext/crates/fp/src/vector/impl_fqvector.rs +++ b/ext/crates/fp/src/vector/impl_fqvector.rs @@ -2,14 +2,10 @@ use std::io; use itertools::Itertools; -use super::{ - inner::{FqSlice, FqSliceMut, FqVector}, - iter::{FqVectorIterator, FqVectorNonZeroIterator}, -}; +use super::inner::FqVector; use crate::{ field::{Field, element::FieldElement}, limb::Limb, - prime::{Prime, ValidPrime}, }; impl FqVector { @@ -51,81 +47,6 @@ impl FqVector { crate::limb::to_bytes(&self.limbs()[..num_limbs], buffer) } - pub const fn is_empty(&self) -> bool { - self.len() == 0 - } - - pub fn prime(&self) -> ValidPrime { - self.fq().characteristic().to_dyn() - } - - #[must_use] - pub fn slice(&self, start: usize, end: usize) -> FqSlice<'_, F> { - assert!(start <= end && end <= self.len()); - FqSlice::new(self.fq(), self.limbs(), start, end) - } - - #[must_use] - pub fn slice_mut(&mut self, start: usize, end: usize) -> FqSliceMut<'_, F> { - assert!(start <= end && end <= self.len()); - FqSliceMut::new(self.fq(), self.limbs_mut(), start, end) - } - - #[inline] - #[must_use] - pub fn as_slice(&self) -> FqSlice<'_, F> { - self.into() - } - - #[inline] - #[must_use] - pub fn as_slice_mut(&mut self) -> FqSliceMut<'_, F> { - self.into() - } - - pub fn add_basis_element(&mut self, index: usize, value: FieldElement) { - assert_eq!(self.fq(), value.field()); - self.as_slice_mut().add_basis_element(index, value); - } - - pub fn entry(&self, index: usize) -> FieldElement { - self.as_slice().entry(index) - } - - pub fn set_entry(&mut self, index: usize, value: FieldElement) { - assert_eq!(self.fq(), value.field()); - self.as_slice_mut().set_entry(index, value); - } - - pub fn iter(&self) -> FqVectorIterator<'_, F> { - self.as_slice().iter() - } - - pub fn iter_nonzero(&self) -> FqVectorNonZeroIterator<'_, F> { - self.as_slice().iter_nonzero() - } - - pub fn set_to_zero(&mut self) { - // This is sound because `fq.encode(fq.zero())` is always zero. - for limb in self.limbs_mut() { - *limb = 0; - } - } - - pub fn scale(&mut self, c: FieldElement) { - assert_eq!(self.fq(), c.field()); - let fq = self.fq(); - - if c == fq.zero() { - self.set_to_zero(); - } - if fq.q() != 2 { - for limb in self.limbs_mut() { - *limb = fq.reduce(fq.fma_limb(0, *limb, c.clone())); - } - } - } - /// Add `other` to `self` on the assumption that the first `offset` entries of `other` are /// empty. pub fn add_offset(&mut self, other: &Self, c: FieldElement, offset: usize) { @@ -176,10 +97,6 @@ impl FqVector { } } - pub fn is_zero(&self) -> bool { - self.limbs().iter().all(|&x| x == 0) - } - /// This function ensures the length of the vector is at least `len`. See also /// `set_scratch_vector_size`. pub fn extend_len(&mut self, len: usize) { @@ -343,12 +260,6 @@ impl From<&FqVector> for Vec> { } } -impl std::fmt::Display for FqVector { - fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - self.as_slice().fmt(f) - } -} - #[cfg(feature = "proptest")] pub mod arbitrary { use proptest::prelude::*; diff --git a/ext/crates/fp/src/vector/inner.rs b/ext/crates/fp/src/vector/inner.rs index 430e0b83e..745352bcc 100644 --- a/ext/crates/fp/src/vector/inner.rs +++ b/ext/crates/fp/src/vector/inner.rs @@ -1,140 +1,394 @@ // This generates better llvm optimization #![allow(clippy::int_plus_one)] -use crate::{field::Field, limb::Limb}; +use std::{ + borrow::Cow, + ops::{Deref, DerefMut, Range}, +}; -/// A vector over a finite field. -/// -/// Interally, it packs entries of the vectors into limbs. However, this is an abstraction that must -/// not leave the `fp` library. -#[derive(Debug, Hash, Eq, PartialEq, Clone)] -pub struct FqVector { - fq: F, - len: usize, - limbs: Vec, -} +use itertools::Itertools; -/// A slice of an `FqVector`. -/// -/// This immutably borrows the vector and implements `Copy`. -#[derive(Debug, Copy, Clone)] -pub struct FqSlice<'a, F: Field> { - fq: F, - limbs: &'a [Limb], - start: usize, - end: usize, -} +use super::iter::{FqVectorIterator, FqVectorNonZeroIterator}; +use crate::{ + field::{Field, element::FieldElement}, + limb::Limb, + prime::{Prime, ValidPrime}, +}; + +pub trait Repr: Deref {} -/// A mutable slice of an `FqVector`. +impl> Repr for T {} + +pub trait ReprMut: DerefMut {} + +impl> ReprMut for T {} + +/// A vector over a finite field. +/// +/// Internally, it packs entries of the vectors into limbs. However, this is an abstraction that +/// must not leave the `fp` library. /// -/// This mutably borrows the vector. Since it is a mutable borrow, it cannot implement `Copy`. -/// However, it has a [`FqSliceMut::copy`] function that imitates the reborrowing, that mutably -/// borrows `FqSliceMut` and returns a `FqSliceMut` with a shorter lifetime. -#[derive(Debug)] -pub struct FqSliceMut<'a, F: Field> { +/// We are generic over a number of types to provide maximal flexibility: +/// - `A` determines whether the vector type is aligned, i.e. if it always starts on a limb +/// boundary. This allows a number of methods to use a fast path. +/// - `R` determines where the limbs are stored. An owned vector will own its limbs in a `Vec`, but +/// a (mutable) slice will hold a (mutable) reference, etc. This allows for more exotic storage +/// options such as `Cow` or `Arc`. +/// - `F` is the underlying field. +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] +pub struct FqVectorBase { fq: F, - limbs: &'a mut [Limb], + limbs: R, start: usize, end: usize, } -// See impl_* for implementations +pub type FqVector = FqVectorBase, F>; +pub type FqSlice<'a, F> = FqVectorBase; +pub type FqSliceMut<'a, F> = FqVectorBase; +pub type FqCow<'a, F> = FqVectorBase, F>; // Accessors -impl FqVector { - pub fn from_raw_parts(fq: F, len: usize, limbs: Vec) -> Self { - debug_assert_eq!(limbs.len(), fq.number(len)); - Self { fq, len, limbs } +impl FqVectorBase { + pub(super) fn _new(fq: F, limbs: R, start: usize, end: usize) -> Self { + assert!(start <= end); + if A { + assert!(start.is_multiple_of(fq.entries_per_limb())); + } + + Self { + fq, + limbs, + start, + end, + } } pub fn fq(&self) -> F { self.fq } - pub const fn len(&self) -> usize { - self.len + pub fn prime(&self) -> ValidPrime { + self.fq().characteristic().to_dyn() } - pub(super) fn limbs(&self) -> &[Limb] { - &self.limbs + pub fn len(&self) -> usize { + self.end() - self.start() } - pub(super) fn limbs_mut(&mut self) -> &mut [Limb] { - &mut self.limbs + pub fn is_empty(&self) -> bool { + self.len() == 0 } - pub(super) fn vec_mut(&mut self) -> &mut Vec { - &mut self.limbs + #[must_use] + pub fn slice(&self, start: usize, end: usize) -> FqSlice<'_, F> { + assert!(start <= end && end <= self.len()); + + FqSlice::_new( + self.fq(), + self.limbs(), + self.start() + start, + self.start() + end, + ) } - pub(super) fn len_mut(&mut self) -> &mut usize { - &mut self.len + pub fn as_slice(&self) -> FqSlice<'_, F> { + self.slice(0, self.len()) } -} -impl<'a, F: Field> FqSlice<'a, F> { - pub(super) fn new(fq: F, limbs: &'a [Limb], start: usize, end: usize) -> Self { - Self { - fq, - limbs, - start, - end, - } + /// TODO: implement prime 2 version + pub fn iter(&self) -> FqVectorIterator<'_, F> { + FqVectorIterator::new(self.as_slice()) } - pub fn fq(&self) -> F { - self.fq + pub fn iter_nonzero(&self) -> FqVectorNonZeroIterator<'_, F> { + FqVectorNonZeroIterator::new(self.as_slice()) } - pub(super) fn into_limbs(self) -> &'a [Limb] { - self.limbs + pub fn is_zero(&self) -> bool { + if A { + return self.limbs().iter().all(|&x| x == 0); + } + + let limb_range = self.limb_range(); + if limb_range.is_empty() { + return true; + } + let (min_mask, max_mask) = self.limb_masks(); + if self.limbs()[limb_range.start] & min_mask != 0 { + return false; + } + + let inner_range = self.limb_range_inner(); + if !inner_range.is_empty() && self.limbs()[inner_range].iter().any(|&x| x != 0) { + return false; + } + if self.limbs()[limb_range.end - 1] & max_mask != 0 { + return false; + } + true + } + + pub fn entry(&self, index: usize) -> FieldElement { + debug_assert!( + index < self.len(), + "Index {} too large, length of vector is only {}.", + index, + self.len() + ); + let bit_mask = self.fq().bitmask(); + let limb_index = self.fq().limb_bit_index_pair(index + self.start()); + let mut result = self.limbs()[limb_index.limb]; + result >>= limb_index.bit_index; + result &= bit_mask; + self.fq().decode(result) } - pub(super) const fn start(&self) -> usize { + // Repr accessors + + pub(super) fn start(&self) -> usize { self.start } - pub(super) const fn end(&self) -> usize { + pub(super) fn end(&self) -> usize { self.end } pub(super) fn limbs(&self) -> &[Limb] { - self.limbs + &self.limbs + } + + // Limb methods + + #[inline] + pub(super) fn offset(&self) -> usize { + let bit_length = self.fq().bit_length(); + let entries_per_limb = self.fq().entries_per_limb(); + (self.start() % entries_per_limb) * bit_length + } + + #[inline] + pub(super) fn limb_range(&self) -> Range { + self.fq().range(self.start(), self.end()) + } + + /// This function underflows if `self.end() == 0`, which happens if and only if we are taking a + /// slice of width 0 at the start of an `FpVector`. This should be a very rare edge case. + /// Dealing with the underflow properly would probably require using `saturating_sub` or + /// something of that nature, and that has a nontrivial (10%) performance hit. + #[inline] + pub(super) fn limb_range_inner(&self) -> Range { + let range = self.limb_range(); + (range.start + 1)..(usize::max(range.start + 1, range.end - 1)) + } + + #[inline(always)] + pub(super) fn min_limb_mask(&self) -> Limb { + !0 << self.offset() + } + + #[inline(always)] + pub(super) fn max_limb_mask(&self) -> Limb { + let num_entries = 1 + (self.end() - 1) % self.fq().entries_per_limb(); + let bit_max = num_entries * self.fq().bit_length(); + + (!0) >> (crate::constants::BITS_PER_LIMB - bit_max) + } + + #[inline(always)] + pub(super) fn limb_masks(&self) -> (Limb, Limb) { + if self.limb_range().len() == 1 { + ( + self.min_limb_mask() & self.max_limb_mask(), + self.min_limb_mask() & self.max_limb_mask(), + ) + } else { + (self.min_limb_mask(), self.max_limb_mask()) + } } } -impl<'a, F: Field> FqSliceMut<'a, F> { - pub(super) fn new(fq: F, limbs: &'a mut [Limb], start: usize, end: usize) -> Self { - Self { - fq, - limbs, - start, - end, +impl FqVectorBase { + #[inline] + #[must_use] + pub fn as_slice_mut(&mut self) -> FqSliceMut<'_, F> { + self.slice_mut(0, self.len()) + } + + pub fn set_to_zero(&mut self) { + if A { + // This is sound because `fq.encode(fq.zero())` is always zero. + for limb in self.limbs_mut() { + *limb = 0; + } + return; + } + + let limb_range = self.limb_range(); + if limb_range.is_empty() { + return; + } + let (min_mask, max_mask) = self.limb_masks(); + self.limbs_mut()[limb_range.start] &= !min_mask; + + let inner_range = self.limb_range_inner(); + for limb in self.limbs_mut()[inner_range].iter_mut() { + *limb = 0; } + self.limbs_mut()[limb_range.end - 1] &= !max_mask; } - pub fn fq(&self) -> F { - self.fq + pub fn scale(&mut self, c: FieldElement) { + assert_eq!(self.fq(), c.field()); + let fq = self.fq(); + + if c == fq.zero() { + self.set_to_zero(); + return; + } + + if fq.q() == 2 { + return; + } + + if A { + for limb in self.limbs_mut() { + *limb = fq.fma_limb(0, *limb, c.clone()); + } + } else { + let limb_range = self.limb_range(); + if limb_range.is_empty() { + return; + } + let (min_mask, max_mask) = self.limb_masks(); + + let limb = self.limbs()[limb_range.start]; + let masked_limb = limb & min_mask; + let rest_limb = limb & !min_mask; + self.limbs_mut()[limb_range.start] = fq.fma_limb(0, masked_limb, c.clone()) | rest_limb; + + let inner_range = self.limb_range_inner(); + for limb in self.limbs_mut()[inner_range].iter_mut() { + *limb = fq.fma_limb(0, *limb, c.clone()); + } + if limb_range.len() > 1 { + let full_limb = self.limbs()[limb_range.end - 1]; + let masked_limb = full_limb & max_mask; + let rest_limb = full_limb & !max_mask; + self.limbs_mut()[limb_range.end - 1] = fq.fma_limb(0, masked_limb, c) | rest_limb; + } + } + + self.reduce_limbs(); } - pub(super) fn start(&self) -> usize { - self.start + pub fn set_entry(&mut self, index: usize, value: FieldElement) { + assert_eq!(self.fq(), value.field()); + assert!(index < self.len()); + + let bit_mask = self.fq().bitmask(); + let limb_index = self.fq().limb_bit_index_pair(index + self.start()); + + let mut result = self.limbs()[limb_index.limb]; + result &= !(bit_mask << limb_index.bit_index); + result |= self.fq().encode(value) << limb_index.bit_index; + self.limbs_mut()[limb_index.limb] = result; } - pub(super) fn end(&self) -> usize { - self.end + pub fn add_basis_element(&mut self, index: usize, value: FieldElement) { + assert_eq!(self.fq(), value.field()); + if self.fq().q() == 2 { + let pair = self.fq().limb_bit_index_pair(index + self.start()); + self.limbs_mut()[pair.limb] ^= self.fq().encode(value) << pair.bit_index; + } else { + let mut x = self.entry(index); + x += value; + self.set_entry(index, x); + } } - pub(super) fn end_mut(&mut self) -> &mut usize { + #[must_use] + pub fn slice_mut(&mut self, start: usize, end: usize) -> FqSliceMut<'_, F> { + assert!(start <= end && end <= self.len()); + let orig_start = self.start(); + + FqSliceMut::_new( + self.fq(), + self.limbs_mut(), + orig_start + start, + orig_start + end, + ) + } + + pub(super) fn limbs_mut(&mut self) -> &mut [Limb] { + &mut self.limbs + } + + pub(super) fn reduce_limbs(&mut self) { + let fq = self.fq(); + if fq.q() != 2 { + let limb_range = self.limb_range(); + + for limb in self.limbs_mut()[limb_range].iter_mut() { + *limb = fq.reduce(*limb); + } + } + } +} + +impl std::fmt::Display for FqVectorBase { + /// # Example + /// ``` + /// # use fp::field::{Field, SmallFq}; + /// # use fp::prime::{P2, ValidPrime}; + /// # use fp::vector::FqVector; + /// let fq = SmallFq::new(P2, 3); + /// let v = FqVector::from_slice(fq, &[fq.zero(), fq.one(), fq.a(), fq.a() * fq.a()]); + /// assert_eq!(&format!("{v}"), "[0, 1, a, a^2]"); + /// + /// // This only looks reasonable over prime fields of order less than 10 + /// assert_eq!(&format!("{v:#}"), "01aa^2"); + /// ``` + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + if f.alternate() { + for v in self.iter() { + // If self.p >= 11, this will look funky + write!(f, "{v}")?; + } + Ok(()) + } else { + write!(f, "[{}]", self.iter().format(", ")) + } + } +} + +// Accessors + +impl FqVector { + pub fn from_raw_parts(fq: F, len: usize, limbs: Vec) -> Self { + debug_assert_eq!(limbs.len(), fq.number(len)); + + Self::_new(fq, limbs, 0, len) + } + + pub(super) fn vec_mut(&mut self) -> &mut Vec { + &mut self.limbs + } + + pub(super) fn len_mut(&mut self) -> &mut usize { &mut self.end } +} - pub(super) fn limbs(&self) -> &[Limb] { +impl<'a, F: Field> FqSlice<'a, F> { + pub(super) fn into_limbs(self) -> &'a [Limb] { self.limbs } +} - pub(super) fn limbs_mut(&mut self) -> &mut [Limb] { - self.limbs +impl<'a, F: Field> FqSliceMut<'a, F> { + pub(super) fn end_mut(&mut self) -> &mut usize { + &mut self.end } }