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::znx_base::{ZnxInfos, ZnxView};
 3use crate::{Backend, DataView, DataViewMut, FFT64, Module, alloc_aligned};
 4use std::marker::PhantomData;
 5
 6const VEC_ZNX_BIG_ROWS: usize = 1;
 7
 8/// VecZnxBig is `Backend` dependent, denoted with backend generic `B`
 9pub struct VecZnxBig<D, B> {
10    data: D,
11    n: usize,
12    cols: usize,
13    size: usize,
14    _phantom: PhantomData<B>,
15}
16
17impl<D, B> ZnxInfos for VecZnxBig<D, B> {
18    fn cols(&self) -> usize {
19        self.cols
20    }
21
22    fn rows(&self) -> usize {
23        1
24    }
25
26    fn n(&self) -> usize {
27        self.n
28    }
29
30    fn size(&self) -> usize {
31        self.size
32    }
33
34    fn sl(&self) -> usize {
35        self.cols() * self.n()
36    }
37}
38
39impl<D, B> DataView for VecZnxBig<D, B> {
40    type D = D;
41    fn data(&self) -> &Self::D {
42        &self.data
43    }
44}
45
46impl<D, B> DataViewMut for VecZnxBig<D, B> {
47    fn data_mut(&mut self) -> &mut Self::D {
48        &mut self.data
49    }
50}
51
52impl<D: AsRef<[u8]>> ZnxView for VecZnxBig<D, FFT64> {
53    type Scalar = i64;
54}
55
56impl<D: From<Vec<u8>>, B: Backend> VecZnxBig<D, B> {
57    pub(crate) fn bytes_of(module: &Module<B>, cols: usize, size: usize) -> usize {
58        unsafe { vec_znx_big::bytes_of_vec_znx_big(module.ptr, size as u64) as usize * cols }
59    }
60
61    pub(crate) fn new(module: &Module<B>, cols: usize, size: usize) -> Self {
62        let data = alloc_aligned::<u8>(Self::bytes_of(module, cols, size));
63        Self {
64            data: data.into(),
65            n: module.n(),
66            cols,
67            size,
68            _phantom: PhantomData,
69        }
70    }
71
72    pub(crate) fn new_from_bytes(module: &Module<B>, cols: usize, size: usize, bytes: impl Into<Vec<u8>>) -> Self {
73        let data: Vec<u8> = bytes.into();
74        assert!(data.len() == Self::bytes_of(module, cols, size));
75        Self {
76            data: data.into(),
77            n: module.n(),
78            cols,
79            size,
80            _phantom: PhantomData,
81        }
82    }
83}
84
85pub type VecZnxBigOwned<B> = VecZnxBig<Vec<u8>, B>;
86
87// impl VecZnxBig<FFT64> {
88//     pub fn print(&self, n: usize, col: usize) {
89//         (0..self.size()).for_each(|i| println!("{}: {:?}", i, &self.at(col, i)[..n]));
90//     }
91// }