poulpy/base2k/src/vec_znx_big.rs
base2k/src/vec_znx_big.rsRaw
1use crate::ffi::vec_znx_big;
2use crate::znx_base::{GetZnxBase, ZnxAlloc, ZnxBase, 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(&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// }