use std::{
fmt,
hash::{DefaultHasher, Hasher},
marker::PhantomData,
};
use crate::{
alloc_aligned,
layouts::{
Backend, Data, DataView, DataViewMut, DigestU64, FillUniform, HostDataMut, HostDataRef, ReaderFrom, ScalarZnx,
ToOwnedDeep, VecZnxInfos, WriterTo, ZnxInfos, ZnxView, ZnxViewMut, ZnxWord, ZnxZero,
},
source::Source,
};
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use rand::Rng;
#[repr(C)]
#[derive(PartialEq, Eq, Clone, Copy, Hash, Debug, Default)]
pub struct VecZnxShape {
n: usize,
cols: usize,
size: usize,
}
impl VecZnxShape {
pub const fn new(n: usize, cols: usize, size: usize) -> Self {
Self { n, cols, size }
}
pub const fn n(self) -> usize {
self.n
}
pub const fn cols(self) -> usize {
self.cols
}
pub const fn size(self) -> usize {
self.size
}
pub(crate) const fn with_size(self, size: usize) -> Self {
assert!(size <= self.size);
Self { size, ..self }
}
}
#[repr(C)]
#[derive(PartialEq, Eq, Clone, Copy, Hash)]
pub struct VecZnx<D: Data, W: ZnxWord> {
pub data: D,
shape: VecZnxShape,
pub _phantom: PhantomData<W>,
}
impl<D: HostDataRef, W: ZnxWord> VecZnx<D, W> {
pub fn as_scalar_znx_ref(&self, col: usize, limb: usize) -> ScalarZnx<&[u8], W> {
ScalarZnx::from_data(bytemuck::cast_slice(self.at(col, limb)), self.n(), 1)
}
}
impl<D: HostDataMut, W: ZnxWord> VecZnx<D, W> {
pub fn as_scalar_znx_mut(&mut self, col: usize, limb: usize) -> ScalarZnx<&mut [u8], W> {
let n = self.n();
ScalarZnx::from_data(bytemuck::cast_slice_mut(self.at_mut(col, limb)), n, 1)
}
}
impl<D: Data + Default, W: ZnxWord> Default for VecZnx<D, W> {
fn default() -> Self {
Self {
data: D::default(),
shape: VecZnxShape::default(),
_phantom: PhantomData,
}
}
}
impl<D: HostDataRef, W: ZnxWord> DigestU64 for VecZnx<D, W> {
fn digest_u64(&self) -> u64 {
let mut h: DefaultHasher = DefaultHasher::new();
h.write(self.data.as_ref());
h.write_usize(self.n());
h.write_usize(self.cols());
h.write_usize(self.size());
h.finish()
}
}
impl<D: HostDataRef, W: ZnxWord> ToOwnedDeep for VecZnx<D, W> {
type Owned = VecZnx<Vec<u8>, W>;
fn to_owned_deep(&self) -> Self::Owned {
VecZnx {
data: self.data.as_ref().to_vec(),
shape: self.shape,
_phantom: PhantomData,
}
}
}
impl<D: Data, W: ZnxWord> VecZnx<D, W> {
pub fn to_host_owned<BE>(&self) -> VecZnx<Vec<u8>, W>
where
BE: Backend<OwnedBuf = D>,
{
let shape = self.shape();
VecZnx::from_data(
crate::layouts::HostBytesBackend::from_bytes(BE::to_host_bytes(&self.data)),
shape.n(),
shape.cols(),
shape.size(),
)
}
pub fn display_host<BE>(&self) -> String
where
BE: Backend<OwnedBuf = D>,
{
self.to_host_owned::<BE>().to_string()
}
}
impl<D: HostDataRef, W: ZnxWord> fmt::Debug for VecZnx<D, W> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{self}")
}
}
impl<D: Data, W: ZnxWord> ZnxInfos for VecZnx<D, W> {
fn n(&self) -> usize {
self.shape.n()
}
fn size(&self) -> usize {
self.shape.size()
}
fn poly_count(&self) -> usize {
crate::layouts::checked_product(&[self.cols(), self.size()], "polynomial count")
}
}
impl<D: Data, W: ZnxWord> VecZnxInfos for VecZnx<D, W> {
fn cols(&self) -> usize {
self.shape.cols()
}
}
impl<D: Data, W: ZnxWord> DataView for VecZnx<D, W> {
type D = D;
fn data(&self) -> &Self::D {
&self.data
}
}
impl<D: Data, W: ZnxWord> DataViewMut for VecZnx<D, W> {
fn data_mut(&mut self) -> &mut Self::D {
&mut self.data
}
}
impl<D: HostDataRef, W: ZnxWord> ZnxView for VecZnx<D, W> {
type Scalar = W;
}
impl<D: Data, W: ZnxWord> VecZnx<D, W> {
pub fn n(&self) -> usize {
self.shape.n()
}
pub fn cols(&self) -> usize {
self.shape.cols()
}
pub fn size(&self) -> usize {
self.shape.size()
}
pub fn shape(&self) -> VecZnxShape {
self.shape
}
}
impl<D: Data, W: ZnxWord> VecZnx<D, W> {
pub fn rsh_tmp_bytes(n: usize) -> usize {
n * size_of::<W>()
}
}
impl<D: HostDataMut, W: ZnxWord> ZnxZero for VecZnx<D, W> {
fn zero(&mut self) {
self.raw_mut().fill(W::zero())
}
fn zero_at(&mut self, i: usize, j: usize) {
self.at_mut(i, j).fill(W::zero());
}
}
impl<D: Data, W: ZnxWord> VecZnx<D, W> {
pub fn bytes_of(n: usize, cols: usize, size: usize) -> usize {
crate::layouts::checked_product(&[n, cols, size, size_of::<W>()], "VecZnx byte size")
}
}
impl<W: ZnxWord> VecZnx<Vec<u8>, W> {
pub(crate) fn alloc(n: usize, cols: usize, size: usize) -> Self {
let data: Vec<u8> = alloc_aligned::<u8>(Self::bytes_of(n, cols, size));
Self {
data,
shape: VecZnxShape::new(n, cols, size),
_phantom: PhantomData,
}
}
pub fn from_bytes(n: usize, cols: usize, size: usize, bytes: impl Into<Vec<u8>>) -> Self {
let data: Vec<u8> = bytes.into();
assert!(
data.len() == Self::bytes_of(n, cols, size),
"from_bytes: data.len()={} != bytes_of({}, {}, {})={}",
data.len(),
n,
cols,
size,
Self::bytes_of(n, cols, size)
);
crate::assert_alignment(data.as_ptr());
Self {
data,
shape: VecZnxShape::new(n, cols, size),
_phantom: PhantomData,
}
}
}
impl<D: Data, W: ZnxWord> VecZnx<D, W> {
pub fn from_data(data: D, n: usize, cols: usize, size: usize) -> Self {
Self {
data,
shape: VecZnxShape::new(n, cols, size),
_phantom: PhantomData,
}
}
}
impl<D: HostDataRef, W: ZnxWord> fmt::Display for VecZnx<D, W> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
writeln!(f, "VecZnx(n={}, cols={}, size={})", self.n(), self.cols(), self.size())?;
for col in 0..self.cols() {
writeln!(f, "Column {col}:")?;
for size in 0..self.size() {
let coeffs = self.at(col, size);
write!(f, " Size {size}: [")?;
let max_show = 16;
let show_count = coeffs.len().min(max_show);
for (i, &coeff) in coeffs.iter().take(show_count).enumerate() {
if i > 0 {
write!(f, ", ")?;
}
write!(f, "{coeff}")?;
}
if coeffs.len() > max_show {
write!(f, ", ... ({} more)", coeffs.len() - max_show)?;
}
writeln!(f, "]")?;
}
}
Ok(())
}
}
impl<D: HostDataMut, W: ZnxWord> FillUniform for VecZnx<D, W> {
fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
assert!(log_bound != 0, "invalid log_bound, cannot be zero");
assert!(
log_bound <= W::BITS,
"log_bound {log_bound} exceeds the {}-bit coefficient word",
W::BITS
);
if log_bound == W::BITS {
source.fill_bytes(self.data.as_mut());
return;
}
let mask: u64 = (1u64 << log_bound) - 1;
let shift: usize = 64 - log_bound;
for x in self.raw_mut().iter_mut() {
let r = source.next_u64() & mask;
*x = W::from_i64(((r << shift) as i64) >> shift);
}
}
}
pub type VecZnxOwned<W> = VecZnx<Vec<u8>, W>;
pub type VecZnxMut<'a, W> = VecZnx<&'a mut [u8], W>;
pub type VecZnxRef<'a, W> = VecZnx<&'a [u8], W>;
pub type VecZnxBackendRef<'a, B> = VecZnx<<B as Backend>::BufRef<'a>, <B as Backend>::ZnxWord>;
pub type VecZnxBackendMut<'a, B> = VecZnx<<B as Backend>::BufMut<'a>, <B as Backend>::ZnxWord>;
pub trait VecZnxAsScalarBackendRef<B: Backend> {
fn as_scalar_znx_backend_ref(&self, col: usize, limb: usize) -> ScalarZnx<B::BufRef<'_>, B::ZnxWord>;
}
impl<B: Backend> VecZnxAsScalarBackendRef<B> for VecZnx<B::OwnedBuf, B::ZnxWord> {
fn as_scalar_znx_backend_ref(&self, col: usize, limb: usize) -> ScalarZnx<B::BufRef<'_>, B::ZnxWord> {
assert!(limb < self.size(), "size: {limb} >= {}", self.size());
assert!(col < self.cols(), "cols: {col} >= {}", self.cols());
let start: usize = limb
.checked_mul(self.cols())
.and_then(|x| x.checked_add(col))
.and_then(|x| x.checked_mul(self.n()))
.and_then(|x| x.checked_mul(B::size_of_znx_word()))
.expect("VecZnx scalar backend view offset overflows usize");
let len: usize = self
.n()
.checked_mul(B::size_of_znx_word())
.expect("VecZnx scalar backend view length overflows usize");
ScalarZnx::from_data(B::region(&self.data, start, len), self.n(), 1)
}
}
pub trait VecZnxAsScalarBackendMut<B: Backend> {
fn as_scalar_znx_backend_mut(&mut self, col: usize, limb: usize) -> ScalarZnx<B::BufMut<'_>, B::ZnxWord>;
}
impl<B: Backend> VecZnxAsScalarBackendMut<B> for VecZnx<B::OwnedBuf, B::ZnxWord> {
fn as_scalar_znx_backend_mut(&mut self, col: usize, limb: usize) -> ScalarZnx<B::BufMut<'_>, B::ZnxWord> {
let n = self.n();
assert!(limb < self.size(), "size: {limb} >= {}", self.size());
assert!(col < self.cols(), "cols: {col} >= {}", self.cols());
let start: usize = limb
.checked_mul(self.cols())
.and_then(|x| x.checked_add(col))
.and_then(|x| x.checked_mul(n))
.and_then(|x| x.checked_mul(B::size_of_znx_word()))
.expect("VecZnx scalar backend view offset overflows usize");
let len: usize = n
.checked_mul(B::size_of_znx_word())
.expect("VecZnx scalar backend view length overflows usize");
ScalarZnx::from_data(B::region_mut(&mut self.data, start, len), n, 1)
}
}
pub trait VecZnxToBackendRef<B: Backend = crate::layouts::HostBytesBackend> {
fn to_backend_ref(&self) -> VecZnxBackendRef<'_, B>;
}
impl<B: Backend> VecZnxToBackendRef<B> for VecZnx<B::OwnedBuf, B::ZnxWord> {
fn to_backend_ref(&self) -> VecZnxBackendRef<'_, B> {
VecZnx {
data: B::view(&self.data),
shape: self.shape,
_phantom: PhantomData,
}
}
}
impl<'b, B: Backend + 'b> VecZnxToBackendRef<B> for &VecZnx<B::BufRef<'b>, B::ZnxWord> {
fn to_backend_ref(&self) -> VecZnxBackendRef<'_, B> {
vec_znx_backend_ref_from_ref::<B>(self)
}
}
impl VecZnxToBackendRef<crate::layouts::HostBytesBackend> for VecZnx<&mut [u8], i64> {
fn to_backend_ref(&self) -> VecZnxBackendRef<'_, crate::layouts::HostBytesBackend> {
VecZnx {
data: self.data,
shape: self.shape,
_phantom: PhantomData,
}
}
}
impl VecZnxToBackendRef<crate::layouts::HostBytesBackend> for VecZnx<&[u8], i64> {
fn to_backend_ref(&self) -> VecZnxBackendRef<'_, crate::layouts::HostBytesBackend> {
VecZnx {
data: self.data,
shape: self.shape,
_phantom: PhantomData,
}
}
}
pub trait VecZnxReborrowBackendRef<B: Backend = crate::layouts::HostBytesBackend> {
fn reborrow_backend_ref(&self) -> VecZnxBackendRef<'_, B>;
}
pub fn vec_znx_backend_ref_from_ref<'a, 'b, B: Backend + 'b>(
vec: &'a VecZnx<B::BufRef<'b>, B::ZnxWord>,
) -> VecZnxBackendRef<'a, B> {
VecZnx {
data: B::view_ref(&vec.data),
shape: vec.shape,
_phantom: PhantomData,
}
}
pub fn vec_znx_backend_ref_from_mut<'a, 'b, B: Backend + 'b>(
vec: &'a VecZnx<B::BufMut<'b>, B::ZnxWord>,
) -> VecZnxBackendRef<'a, B> {
VecZnx {
data: B::view_ref_mut(&vec.data),
shape: vec.shape,
_phantom: PhantomData,
}
}
impl<'b, B: Backend + 'b> VecZnxReborrowBackendRef<B> for VecZnx<B::BufMut<'b>, B::ZnxWord> {
fn reborrow_backend_ref(&self) -> VecZnxBackendRef<'_, B> {
vec_znx_backend_ref_from_mut::<B>(self)
}
}
pub trait VecZnxToBackendMut<B: Backend = crate::layouts::HostBytesBackend> {
fn to_backend_mut(&mut self) -> VecZnxBackendMut<'_, B>;
}
impl<B: Backend> VecZnxToBackendMut<B> for VecZnx<B::OwnedBuf, B::ZnxWord> {
fn to_backend_mut(&mut self) -> VecZnxBackendMut<'_, B> {
VecZnx {
data: B::view_mut(&mut self.data),
shape: self.shape,
_phantom: PhantomData,
}
}
}
impl<'b, B: Backend + 'b> VecZnxToBackendMut<B> for &mut VecZnx<B::BufMut<'b>, B::ZnxWord> {
fn to_backend_mut(&mut self) -> VecZnxBackendMut<'_, B> {
vec_znx_backend_mut_from_mut::<B>(self)
}
}
impl VecZnxToBackendMut<crate::layouts::HostBytesBackend> for VecZnx<&mut [u8], i64> {
fn to_backend_mut(&mut self) -> VecZnxBackendMut<'_, crate::layouts::HostBytesBackend> {
VecZnx {
data: self.data,
shape: self.shape,
_phantom: PhantomData,
}
}
}
pub trait VecZnxReborrowBackendMut<B: Backend = crate::layouts::HostBytesBackend> {
fn reborrow_backend_mut(&mut self) -> VecZnxBackendMut<'_, B>;
}
pub fn vec_znx_host_backend_ref<D: HostDataRef>(vec: &VecZnx<D, i64>) -> VecZnxBackendRef<'_, crate::layouts::HostBytesBackend> {
VecZnx {
data: vec.data.as_ref(),
shape: vec.shape,
_phantom: PhantomData,
}
}
pub fn vec_znx_host_backend_mut<D: HostDataMut>(
vec: &mut VecZnx<D, i64>,
) -> VecZnxBackendMut<'_, crate::layouts::HostBytesBackend> {
VecZnx {
data: vec.data.as_mut(),
shape: vec.shape,
_phantom: PhantomData,
}
}
pub fn vec_znx_backend_mut_from_mut<'a, 'b, B: Backend + 'b>(
vec: &'a mut VecZnx<B::BufMut<'b>, B::ZnxWord>,
) -> VecZnxBackendMut<'a, B> {
VecZnx {
data: B::view_mut_ref(&mut vec.data),
shape: vec.shape,
_phantom: PhantomData,
}
}
impl<'b, B: Backend + 'b> VecZnxReborrowBackendMut<B> for VecZnx<B::BufMut<'b>, B::ZnxWord> {
fn reborrow_backend_mut(&mut self) -> VecZnxBackendMut<'_, B> {
vec_znx_backend_mut_from_mut::<B>(self)
}
}
pub fn vec_znx_backend_ref<'a, B: Backend>(vec: &'a VecZnx<B::OwnedBuf, B::ZnxWord>) -> VecZnxBackendRef<'a, B> {
<VecZnx<B::OwnedBuf, B::ZnxWord> as VecZnxToBackendRef<B>>::to_backend_ref(vec)
}
pub fn vec_znx_reborrow_backend_mut<'a, B: Backend>(vec: &'a mut VecZnxBackendMut<'_, B>) -> VecZnxBackendMut<'a, B> {
<VecZnx<B::BufMut<'_>, B::ZnxWord> as VecZnxReborrowBackendMut<B>>::reborrow_backend_mut(vec)
}
pub fn vec_znx_backend_mut<'a, B: Backend>(vec: &'a mut VecZnx<B::OwnedBuf, B::ZnxWord>) -> VecZnxBackendMut<'a, B> {
<VecZnx<B::OwnedBuf, B::ZnxWord> as VecZnxToBackendMut<B>>::to_backend_mut(vec)
}
pub fn vec_znx_backend_mut_with_size<'a, B: Backend>(vec: VecZnxBackendMut<'a, B>, size: usize) -> VecZnxBackendMut<'a, B> {
VecZnx {
data: vec.data,
shape: vec.shape.with_size(size),
_phantom: PhantomData,
}
}
impl<D: HostDataMut, W: ZnxWord> ReaderFrom for VecZnx<D, W> {
fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
let new_n: usize = reader.read_u64::<LittleEndian>()? as usize;
let new_cols: usize = reader.read_u64::<LittleEndian>()? as usize;
let new_size: usize = reader.read_u64::<LittleEndian>()? as usize;
let len: usize = reader.read_u64::<LittleEndian>()? as usize;
let expected_len: usize =
crate::layouts::checked_product(&[new_n, new_cols, new_size, size_of::<W>()], "VecZnx serialized byte size");
if expected_len != len {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!(
"VecZnx metadata inconsistent: n={new_n} * cols={new_cols} * size={new_size} * {} = {expected_len} != data len={len}",
size_of::<W>()
),
));
}
let buf: &mut [u8] = self.data.as_mut();
if buf.len() < len {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("VecZnx buffer too small: self.data.len()={} < read len={len}", buf.len()),
));
}
reader.read_exact(&mut buf[..len])?;
self.shape = VecZnxShape::new(new_n, new_cols, new_size);
Ok(())
}
}
impl<D: HostDataRef, W: ZnxWord> WriterTo for VecZnx<D, W> {
fn write_to<Wr: std::io::Write>(&self, writer: &mut Wr) -> std::io::Result<()> {
writer.write_u64::<LittleEndian>(self.n() as u64)?;
writer.write_u64::<LittleEndian>(self.cols() as u64)?;
writer.write_u64::<LittleEndian>(self.size() as u64)?;
let coeff_bytes: usize = crate::layouts::checked_product(
&[self.n(), self.cols(), self.size(), size_of::<W>()],
"VecZnx logical byte size",
);
let buf: &[u8] = self.data.as_ref();
if buf.len() < coeff_bytes {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!(
"VecZnx buffer too small: self.data.len()={} < coeff_bytes={coeff_bytes}",
buf.len()
),
));
}
writer.write_u64::<LittleEndian>(coeff_bytes as u64)?;
writer.write_all(&buf[..coeff_bytes])?;
Ok(())
}
}