1use crate::ffi::vec_znx_big;
 2use crate::znx_base::{GetZnxBase, ZnxAlloc, ZnxBase, ZnxInfos, ZnxLayout, ZnxSliceSize, ZnxZero};
 3use crate::{Backend, FFT64, Module, NTT120};
 4use std::marker::PhantomData;
 5
 6const VEC_ZNX_BIG_ROWS: usize = 1;
 7
 8pub struct VecZnxBig<B: Backend> {
 9    pub inner: ZnxBase,
10    pub _marker: PhantomData<B>,
11}
12
13impl<B: Backend> GetZnxBase for VecZnxBig<B> {
14    fn znx(&self) -> &ZnxBase {
15        &self.inner
16    }
17
18    fn znx_mut(&mut self) -> &mut ZnxBase {
19        &mut self.inner
20    }
21}
22
23impl<B: Backend> ZnxInfos for VecZnxBig<B> {}
24
25impl<B: Backend> ZnxAlloc<B> for VecZnxBig<B> {
26    type Scalar = u8;
27
28    fn from_bytes_borrow(module: &Module<B>, _rows: usize, cols: usize, size: usize, bytes: &mut [u8]) -> Self {
29        debug_assert_eq!(bytes.len(), Self::bytes_of(module, _rows, cols, size));
30        VecZnxBig {
31            inner: ZnxBase::from_bytes_borrow(module.n(), VEC_ZNX_BIG_ROWS, cols, size, bytes),
32            _marker: PhantomData,
33        }
34    }
35
36    fn bytes_of(module: &Module<B>, _rows: usize, cols: usize, size: usize) -> usize {
37        debug_assert_eq!(
38            _rows, VEC_ZNX_BIG_ROWS,
39            "rows != {} not supported for VecZnxBig",
40            VEC_ZNX_BIG_ROWS
41        );
42        unsafe { vec_znx_big::bytes_of_vec_znx_big(module.ptr, size as u64) as usize * cols }
43    }
44}
45
46impl ZnxLayout for VecZnxBig<FFT64> {
47    type Scalar = i64;
48}
49
50impl ZnxLayout for VecZnxBig<NTT120> {
51    type Scalar = i128;
52}
53
54impl ZnxZero for VecZnxBig<FFT64> {}
55
56impl ZnxSliceSize for VecZnxBig<FFT64> {
57    fn sl(&self) -> usize {
58        self.n() * self.cols()
59    }
60}
61
62impl ZnxSliceSize for VecZnxBig<NTT120> {
63    fn sl(&self) -> usize {
64        self.n() * 4 * self.cols()
65    }
66}
67
68impl ZnxZero for VecZnxBig<NTT120> {}
69
70impl VecZnxBig<FFT64> {
71    pub fn print(&self, n: usize, col: usize) {
72        (0..self.size()).for_each(|i| println!("{}: {:?}", i, &self.at(col, i)[..n]));
73    }
74}