Add Zn type

poulpy-hal/src/layouts/encoding.rs

@@ -3,7 +3,7 @@ use rug::{Assign, Float};
 
 use crate::{
     api::{ZnxInfos, ZnxView, ZnxViewMut, ZnxZero},
-    layouts::{DataMut, DataRef, VecZnx, VecZnxToMut, VecZnxToRef},
+    layouts::{DataMut, DataRef, VecZnx, VecZnxToMut, VecZnxToRef, Zn, ZnToMut, ZnToRef},
 };
 
 impl<D: DataMut> VecZnx<D> {
@@ -202,3 +202,90 @@ impl<D: DataRef> VecZnx<D> {
         });
     }
 }
+
+impl<D: DataMut> Zn<D> {
+    pub fn encode_i64(&mut self, basek: usize, k: usize, data: i64, log_max: usize) {
+        let size: usize = k.div_ceil(basek);
+
+        #[cfg(debug_assertions)]
+        {
+            let a: Zn<&mut [u8]> = self.to_mut();
+            assert!(
+                size <= a.size(),
+                "invalid argument k.div_ceil(basek)={} > a.size()={}",
+                size,
+                a.size()
+            );
+        }
+
+        let k_rem: usize = basek - (k % basek);
+        let mut a: Zn<&mut [u8]> = self.to_mut();
+        (0..a.size()).for_each(|j| a.at_mut(0, j)[0] = 0);
+
+        // If 2^{basek} * 2^{k_rem} < 2^{63}-1, then we can simply copy
+        // values on the last limb.
+        // Else we decompose values base2k.
+        if log_max + k_rem < 63 || k_rem == basek {
+            a.at_mut(0, size - 1)[0] = data;
+        } else {
+            let mask: i64 = (1 << basek) - 1;
+            let steps: usize = size.min(log_max.div_ceil(basek));
+            (size - steps..size)
+                .rev()
+                .enumerate()
+                .for_each(|(j, j_rev)| {
+                    a.at_mut(0, j_rev)[0] = (data >> (j * basek)) & mask;
+                })
+        }
+
+        // Case where prec % k != 0.
+        if k_rem != basek {
+            let steps: usize = size.min(log_max.div_ceil(basek));
+            (size - steps..size).rev().for_each(|j| {
+                a.at_mut(0, j)[0] <<= k_rem;
+            })
+        }
+    }
+}
+
+impl<D: DataRef> Zn<D> {
+    pub fn decode_i64(&self, basek: usize, k: usize) -> i64 {
+        let a: Zn<&[u8]> = self.to_ref();
+        let size: usize = k.div_ceil(basek);
+        let mut res: i64 = 0;
+        let rem: usize = basek - (k % basek);
+        (0..size).for_each(|j| {
+            let x: i64 = a.at(0, j)[0];
+            if j == size - 1 && rem != basek {
+                let k_rem: usize = basek - rem;
+                res = (res << k_rem) + (x >> rem);
+            } else {
+                res = (res << basek) + x;
+            }
+        });
+        res
+    }
+
+    pub fn decode_float(&self, basek: usize) -> Float {
+        let a: Zn<&[u8]> = self.to_ref();
+        let size: usize = a.size();
+        let prec: u32 = (basek * size) as u32;
+
+        // 2^{basek}
+        let base: Float = Float::with_val(prec, (1 << basek) as f64);
+        let mut res: Float = Float::with_val(prec, (1 << basek) as f64);
+
+        // y[i] = sum x[j][i] * 2^{-basek*j}
+        (0..size).for_each(|i| {
+            if i == 0 {
+                res.assign(a.at(0, size - i - 1)[0]);
+                res /= &base;
+            } else {
+                res += Float::with_val(prec, a.at(0, size - i - 1)[0]);
+                res /= &base;
+            }
+        });
+
+        res
+    }
+}

poulpy-hal/src/layouts/mod.rs

@@ -10,6 +10,7 @@ mod vec_znx;
 mod vec_znx_big;
 mod vec_znx_dft;
 mod vmp_pmat;
+mod zn;
 
 pub use mat_znx::*;
 pub use module::*;
@@ -21,6 +22,7 @@ pub use vec_znx::*;
 pub use vec_znx_big::*;
 pub use vec_znx_dft::*;
 pub use vmp_pmat::*;
+pub use zn::*;
 
 pub trait Data = PartialEq + Eq + Sized;
 pub trait DataRef = Data + AsRef<[u8]>;

poulpy-hal/src/layouts/zn.rs

@@ -0,0 +1,255 @@
+use std::fmt;
+
+use crate::{
+    alloc_aligned,
+    api::{DataView, DataViewMut, FillUniform, Reset, ZnxInfos, ZnxSliceSize, ZnxView, ZnxViewMut, ZnxZero},
+    layouts::{Data, DataMut, DataRef, ReaderFrom, ToOwnedDeep, WriterTo},
+    source::Source,
+};
+
+use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
+use rand::RngCore;
+
+#[derive(PartialEq, Eq, Clone, Copy)]
+pub struct Zn<D: Data> {
+    pub data: D,
+    pub n: usize,
+    pub cols: usize,
+    pub size: usize,
+    pub max_size: usize,
+}
+
+impl<D: DataRef> ToOwnedDeep for Zn<D> {
+    type Owned = Zn<Vec<u8>>;
+    fn to_owned_deep(&self) -> Self::Owned {
+        Zn {
+            data: self.data.as_ref().to_vec(),
+            n: self.n,
+            cols: self.cols,
+            size: self.size,
+            max_size: self.max_size,
+        }
+    }
+}
+
+impl<D: DataRef> fmt::Debug for Zn<D> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", self)
+    }
+}
+
+impl<D: Data> ZnxInfos for Zn<D> {
+    fn cols(&self) -> usize {
+        self.cols
+    }
+
+    fn rows(&self) -> usize {
+        1
+    }
+
+    fn n(&self) -> usize {
+        self.n
+    }
+
+    fn size(&self) -> usize {
+        self.size
+    }
+}
+
+impl<D: Data> ZnxSliceSize for Zn<D> {
+    fn sl(&self) -> usize {
+        self.n() * self.cols()
+    }
+}
+
+impl<D: Data> DataView for Zn<D> {
+    type D = D;
+    fn data(&self) -> &Self::D {
+        &self.data
+    }
+}
+
+impl<D: Data> DataViewMut for Zn<D> {
+    fn data_mut(&mut self) -> &mut Self::D {
+        &mut self.data
+    }
+}
+
+impl<D: DataRef> ZnxView for Zn<D> {
+    type Scalar = i64;
+}
+
+impl Zn<Vec<u8>> {
+    pub fn rsh_scratch_space(n: usize) -> usize {
+        n * std::mem::size_of::<i64>()
+    }
+}
+
+impl<D: DataMut> ZnxZero for Zn<D> {
+    fn zero(&mut self) {
+        self.raw_mut().fill(0)
+    }
+    fn zero_at(&mut self, i: usize, j: usize) {
+        self.at_mut(i, j).fill(0);
+    }
+}
+
+impl Zn<Vec<u8>> {
+    pub fn alloc_bytes(n: usize, cols: usize, size: usize) -> usize {
+        n * cols * size * size_of::<i64>()
+    }
+
+    pub fn alloc(n: usize, cols: usize, size: usize) -> Self {
+        let data: Vec<u8> = alloc_aligned::<u8>(Self::alloc_bytes(n, cols, size));
+        Self {
+            data,
+            n,
+            cols,
+            size,
+            max_size: size,
+        }
+    }
+
+    pub fn from_bytes<Scalar: Sized>(n: usize, cols: usize, size: usize, bytes: impl Into<Vec<u8>>) -> Self {
+        let data: Vec<u8> = bytes.into();
+        assert!(data.len() == Self::alloc_bytes(n, cols, size));
+        Self {
+            data,
+            n,
+            cols,
+            size,
+            max_size: size,
+        }
+    }
+}
+
+impl<D: Data> Zn<D> {
+    pub fn from_data(data: D, n: usize, cols: usize, size: usize) -> Self {
+        Self {
+            data,
+            n,
+            cols,
+            size,
+            max_size: size,
+        }
+    }
+}
+
+impl<D: DataRef> fmt::Display for Zn<D> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        writeln!(
+            f,
+            "Zn(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 = 100;
+                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: DataMut> FillUniform for Zn<D> {
+    fn fill_uniform(&mut self, source: &mut Source) {
+        source.fill_bytes(self.data.as_mut());
+    }
+}
+
+impl<D: DataMut> Reset for Zn<D> {
+    fn reset(&mut self) {
+        self.zero();
+        self.n = 0;
+        self.cols = 0;
+        self.size = 0;
+        self.max_size = 0;
+    }
+}
+
+pub type ZnOwned = Zn<Vec<u8>>;
+pub type ZnMut<'a> = Zn<&'a mut [u8]>;
+pub type ZnRef<'a> = Zn<&'a [u8]>;
+
+pub trait ZnToRef {
+    fn to_ref(&self) -> Zn<&[u8]>;
+}
+
+impl<D: DataRef> ZnToRef for Zn<D> {
+    fn to_ref(&self) -> Zn<&[u8]> {
+        Zn {
+            data: self.data.as_ref(),
+            n: self.n,
+            cols: self.cols,
+            size: self.size,
+            max_size: self.max_size,
+        }
+    }
+}
+
+pub trait ZnToMut {
+    fn to_mut(&mut self) -> Zn<&mut [u8]>;
+}
+
+impl<D: DataMut> ZnToMut for Zn<D> {
+    fn to_mut(&mut self) -> Zn<&mut [u8]> {
+        Zn {
+            data: self.data.as_mut(),
+            n: self.n,
+            cols: self.cols,
+            size: self.size,
+            max_size: self.max_size,
+        }
+    }
+}
+
+impl<D: DataMut> ReaderFrom for Zn<D> {
+    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
+        self.n = reader.read_u64::<LittleEndian>()? as usize;
+        self.cols = reader.read_u64::<LittleEndian>()? as usize;
+        self.size = reader.read_u64::<LittleEndian>()? as usize;
+        self.max_size = reader.read_u64::<LittleEndian>()? as usize;
+        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
+        let buf: &mut [u8] = self.data.as_mut();
+        if buf.len() != len {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::UnexpectedEof,
+                format!("self.data.len()={} != read len={}", buf.len(), len),
+            ));
+        }
+        reader.read_exact(&mut buf[..len])?;
+        Ok(())
+    }
+}
+
+impl<D: DataRef> WriterTo for Zn<D> {
+    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> 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)?;
+        writer.write_u64::<LittleEndian>(self.max_size as u64)?;
+        let buf: &[u8] = self.data.as_ref();
+        writer.write_u64::<LittleEndian>(buf.len() as u64)?;
+        writer.write_all(buf)?;
+        Ok(())
+    }
+}