1use crate::ffi::vec_znx_big;
  2use crate::znx_base::{ZnxInfos, ZnxView};
  3use crate::{Backend, DataView, DataViewMut, FFT64, Module, ZnxSliceSize, alloc_aligned};
  4use std::marker::PhantomData;
  5
  6pub struct VecZnxBig<D, B: Backend> {
  7    data: D,
  8    n: usize,
  9    cols: usize,
 10    size: usize,
 11    _phantom: PhantomData<B>,
 12}
 13
 14impl<D, B: Backend> ZnxInfos for VecZnxBig<D, B> {
 15    fn cols(&self) -> usize {
 16        self.cols
 17    }
 18
 19    fn rows(&self) -> usize {
 20        1
 21    }
 22
 23    fn n(&self) -> usize {
 24        self.n
 25    }
 26
 27    fn size(&self) -> usize {
 28        self.size
 29    }
 30}
 31
 32impl<D> ZnxSliceSize for VecZnxBig<D, FFT64> {
 33    fn sl(&self) -> usize {
 34        self.n() * self.cols()
 35    }
 36}
 37
 38impl<D, B: Backend> DataView for VecZnxBig<D, B> {
 39    type D = D;
 40    fn data(&self) -> &Self::D {
 41        &self.data
 42    }
 43}
 44
 45impl<D, B: Backend> DataViewMut for VecZnxBig<D, B> {
 46    fn data_mut(&mut self) -> &mut Self::D {
 47        &mut self.data
 48    }
 49}
 50
 51impl<D: AsRef<[u8]>> ZnxView for VecZnxBig<D, FFT64> {
 52    type Scalar = i64;
 53}
 54
 55pub(crate) fn bytes_of_vec_znx_big<B: Backend>(module: &Module<B>, cols: usize, size: usize) -> usize {
 56    unsafe { vec_znx_big::bytes_of_vec_znx_big(module.ptr, size as u64) as usize * cols }
 57}
 58
 59impl<D: From<Vec<u8>>, B: Backend> VecZnxBig<D, B> {
 60    pub(crate) fn new(module: &Module<B>, cols: usize, size: usize) -> Self {
 61        let data = alloc_aligned::<u8>(bytes_of_vec_znx_big(module, cols, size));
 62        Self {
 63            data: data.into(),
 64            n: module.n(),
 65            cols,
 66            size,
 67            _phantom: PhantomData,
 68        }
 69    }
 70
 71    pub(crate) fn new_from_bytes(module: &Module<B>, cols: usize, size: usize, bytes: impl Into<Vec<u8>>) -> Self {
 72        let data: Vec<u8> = bytes.into();
 73        assert!(data.len() == bytes_of_vec_znx_big(module, cols, size));
 74        Self {
 75            data: data.into(),
 76            n: module.n(),
 77            cols,
 78            size,
 79            _phantom: PhantomData,
 80        }
 81    }
 82}
 83
 84impl<D, B: Backend> VecZnxBig<D, B> {
 85    pub(crate) fn from_data(data: D, n: usize, cols: usize, size: usize) -> Self {
 86        Self {
 87            data,
 88            n,
 89            cols,
 90            size,
 91            _phantom: PhantomData,
 92        }
 93    }
 94}
 95
 96pub type VecZnxBigOwned<B> = VecZnxBig<Vec<u8>, B>;
 97
 98pub trait VecZnxBigToRef<B: Backend> {
 99    fn to_ref(&self) -> VecZnxBig<&[u8], B>;
100}
101
102pub trait VecZnxBigToMut<B: Backend> {
103    fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B>;
104}
105
106impl<B: Backend> VecZnxBigToMut<B> for VecZnxBig<Vec<u8>, B> {
107    fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B> {
108        VecZnxBig {
109            data: self.data.as_mut_slice(),
110            n: self.n,
111            cols: self.cols,
112            size: self.size,
113            _phantom: PhantomData,
114        }
115    }
116}
117
118impl<B: Backend> VecZnxBigToRef<B> for VecZnxBig<Vec<u8>, B> {
119    fn to_ref(&self) -> VecZnxBig<&[u8], B> {
120        VecZnxBig {
121            data: self.data.as_slice(),
122            n: self.n,
123            cols: self.cols,
124            size: self.size,
125            _phantom: PhantomData,
126        }
127    }
128}
129
130impl<B: Backend> VecZnxBigToMut<B> for VecZnxBig<&mut [u8], B> {
131    fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B> {
132        VecZnxBig {
133            data: self.data,
134            n: self.n,
135            cols: self.cols,
136            size: self.size,
137            _phantom: PhantomData,
138        }
139    }
140}
141
142impl<B: Backend> VecZnxBigToRef<B> for VecZnxBig<&mut [u8], B> {
143    fn to_ref(&self) -> VecZnxBig<&[u8], B> {
144        VecZnxBig {
145            data: self.data,
146            n: self.n,
147            cols: self.cols,
148            size: self.size,
149            _phantom: PhantomData,
150        }
151    }
152}
153
154impl<B: Backend> VecZnxBigToRef<B> for VecZnxBig<&[u8], B> {
155    fn to_ref(&self) -> VecZnxBig<&[u8], B> {
156        VecZnxBig {
157            data: self.data,
158            n: self.n,
159            cols: self.cols,
160            size: self.size,
161            _phantom: PhantomData,
162        }
163    }
164}