arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/base2k/src/vec_znx_big.rs
base2k/src/vec_znx_big.rsRaw
  1use crate::ffi::vec_znx_big;
  2use crate::{Backend, FFT64, Module, ZnxBase, ZnxBasics, ZnxInfos, ZnxLayout, alloc_aligned, assert_alignement};
  3use std::marker::PhantomData;
  4
  5pub struct VecZnxBig<B: Backend> {
  6    pub data: Vec<u8>,
  7    pub ptr: *mut u8,
  8    pub n: usize,
  9    pub cols: usize,
 10    pub size: usize,
 11    pub _marker: PhantomData<B>,
 12}
 13
 14impl ZnxBasics for VecZnxBig<FFT64> {}
 15
 16impl<B: Backend> ZnxBase<B> for VecZnxBig<B> {
 17    type Scalar = u8;
 18
 19    fn new(module: &Module<B>, cols: usize, size: usize) -> Self {
 20        #[cfg(debug_assertions)]
 21        {
 22            assert!(cols > 0);
 23            assert!(size > 0);
 24        }
 25        let mut data: Vec<Self::Scalar> = alloc_aligned(Self::bytes_of(module, cols, size));
 26        let ptr: *mut Self::Scalar = data.as_mut_ptr();
 27        Self {
 28            data: data,
 29            ptr: ptr,
 30            n: module.n(),
 31            cols: cols,
 32            size: size,
 33            _marker: PhantomData,
 34        }
 35    }
 36
 37    fn bytes_of(module: &Module<B>, cols: usize, size: usize) -> usize {
 38        unsafe { vec_znx_big::bytes_of_vec_znx_big(module.ptr, size as u64) as usize * cols }
 39    }
 40
 41    /// Returns a new [VecZnxBig] with the provided data as backing array.
 42    /// User must ensure that data is properly alligned and that
 43    /// the size of data is at least equal to [Module::bytes_of_vec_znx_big].
 44    fn from_bytes(module: &Module<B>, cols: usize, size: usize, bytes: &mut [Self::Scalar]) -> Self {
 45        #[cfg(debug_assertions)]
 46        {
 47            assert!(cols > 0);
 48            assert!(size > 0);
 49            assert_eq!(bytes.len(), Self::bytes_of(module, cols, size));
 50            assert_alignement(bytes.as_ptr())
 51        };
 52        unsafe {
 53            Self {
 54                data: Vec::from_raw_parts(bytes.as_mut_ptr(), bytes.len(), bytes.len()),
 55                ptr: bytes.as_mut_ptr(),
 56                n: module.n(),
 57                cols: cols,
 58                size: size,
 59                _marker: PhantomData,
 60            }
 61        }
 62    }
 63
 64    fn from_bytes_borrow(module: &Module<B>, cols: usize, size: usize, bytes: &mut [Self::Scalar]) -> Self {
 65        #[cfg(debug_assertions)]
 66        {
 67            assert!(cols > 0);
 68            assert!(size > 0);
 69            assert_eq!(bytes.len(), Self::bytes_of(module, cols, size));
 70            assert_alignement(bytes.as_ptr());
 71        }
 72        Self {
 73            data: Vec::new(),
 74            ptr: bytes.as_mut_ptr(),
 75            n: module.n(),
 76            cols: cols,
 77            size: size,
 78            _marker: PhantomData,
 79        }
 80    }
 81}
 82
 83impl<B: Backend> ZnxInfos for VecZnxBig<B> {
 84    fn n(&self) -> usize {
 85        self.n
 86    }
 87
 88    fn cols(&self) -> usize {
 89        self.cols
 90    }
 91
 92    fn rows(&self) -> usize {
 93        1
 94    }
 95
 96    fn size(&self) -> usize {
 97        self.size
 98    }
 99}
100
101impl ZnxLayout for VecZnxBig<FFT64> {
102    type Scalar = i64;
103
104    fn as_ptr(&self) -> *const Self::Scalar {
105        self.ptr as *const Self::Scalar
106    }
107
108    fn as_mut_ptr(&mut self) -> *mut Self::Scalar {
109        self.ptr as *mut Self::Scalar
110    }
111}
112
113impl VecZnxBig<FFT64> {
114    pub fn print(&self, n: usize) {
115        (0..self.size()).for_each(|i| println!("{}: {:?}", i, &self.at_limb(i)[..n]));
116    }
117}