1use crate::ffi::vec_znx_big;
  2use crate::znx_base::{ZnxInfos, ZnxView};
  3use crate::{Backend, DataView, DataViewMut, FFT64, Module, ZnxSliceSize, ZnxViewMut, ZnxZero, alloc_aligned};
  4use std::fmt;
  5use std::marker::PhantomData;
  6
  7pub struct VecZnxBig<D, B: Backend> {
  8    data: D,
  9    n: usize,
 10    cols: usize,
 11    size: usize,
 12    _phantom: PhantomData<B>,
 13}
 14
 15impl<D, B: Backend> ZnxInfos for VecZnxBig<D, B> {
 16    fn cols(&self) -> usize {
 17        self.cols
 18    }
 19
 20    fn rows(&self) -> usize {
 21        1
 22    }
 23
 24    fn n(&self) -> usize {
 25        self.n
 26    }
 27
 28    fn size(&self) -> usize {
 29        self.size
 30    }
 31}
 32
 33impl<D> ZnxSliceSize for VecZnxBig<D, FFT64> {
 34    fn sl(&self) -> usize {
 35        self.n() * self.cols()
 36    }
 37}
 38
 39impl<D, B: Backend> DataView for VecZnxBig<D, B> {
 40    type D = D;
 41    fn data(&self) -> &Self::D {
 42        &self.data
 43    }
 44}
 45
 46impl<D, B: Backend> 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
 56pub(crate) fn bytes_of_vec_znx_big<B: Backend>(module: &Module<B>, cols: usize, size: usize) -> usize {
 57    unsafe { vec_znx_big::bytes_of_vec_znx_big(module.ptr, size as u64) as usize * cols }
 58}
 59
 60impl<D: From<Vec<u8>>, B: Backend> VecZnxBig<D, B> {
 61    pub(crate) fn new(module: &Module<B>, cols: usize, size: usize) -> Self {
 62        let data = alloc_aligned::<u8>(bytes_of_vec_znx_big(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() == bytes_of_vec_znx_big(module, cols, size));
 75        Self {
 76            data: data.into(),
 77            n: module.n(),
 78            cols,
 79            size,
 80            _phantom: PhantomData,
 81        }
 82    }
 83}
 84
 85impl<D, B: Backend> VecZnxBig<D, B> {
 86    pub(crate) fn from_data(data: D, n: usize, cols: usize, size: usize) -> Self {
 87        Self {
 88            data,
 89            n,
 90            cols,
 91            size,
 92            _phantom: PhantomData,
 93        }
 94    }
 95}
 96
 97impl<D> VecZnxBig<D, FFT64>
 98where
 99    VecZnxBig<D, FFT64>: VecZnxBigToMut<FFT64> + ZnxInfos,
100{
101    /// Extracts the a_col-th column of 'a' and stores it on the self_col-th column [Self].
102    pub fn extract_column<C>(&mut self, self_col: usize, a: &VecZnxBig<C, FFT64>, a_col: usize)
103    where
104        VecZnxBig<C, FFT64>: VecZnxBigToRef<FFT64> + ZnxInfos,
105    {
106        #[cfg(debug_assertions)]
107        {
108            assert!(self_col < self.cols());
109            assert!(a_col < a.cols());
110        }
111
112        let min_size: usize = self.size.min(a.size());
113        let max_size: usize = self.size;
114
115        let mut self_mut: VecZnxBig<&mut [u8], FFT64> = self.to_mut();
116        let a_ref: VecZnxBig<&[u8], FFT64> = a.to_ref();
117
118        (0..min_size).for_each(|i: usize| {
119            self_mut
120                .at_mut(self_col, i)
121                .copy_from_slice(a_ref.at(a_col, i));
122        });
123
124        (min_size..max_size).for_each(|i| {
125            self_mut.zero_at(self_col, i);
126        });
127    }
128}
129
130pub type VecZnxBigOwned<B> = VecZnxBig<Vec<u8>, B>;
131
132pub trait VecZnxBigToRef<B: Backend> {
133    fn to_ref(&self) -> VecZnxBig<&[u8], B>;
134}
135
136pub trait VecZnxBigToMut<B: Backend> {
137    fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B>;
138}
139
140impl<B: Backend> VecZnxBigToMut<B> for VecZnxBig<Vec<u8>, B> {
141    fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B> {
142        VecZnxBig {
143            data: self.data.as_mut_slice(),
144            n: self.n,
145            cols: self.cols,
146            size: self.size,
147            _phantom: PhantomData,
148        }
149    }
150}
151
152impl<B: Backend> VecZnxBigToRef<B> for VecZnxBig<Vec<u8>, B> {
153    fn to_ref(&self) -> VecZnxBig<&[u8], B> {
154        VecZnxBig {
155            data: self.data.as_slice(),
156            n: self.n,
157            cols: self.cols,
158            size: self.size,
159            _phantom: PhantomData,
160        }
161    }
162}
163
164impl<B: Backend> VecZnxBigToMut<B> for VecZnxBig<&mut [u8], B> {
165    fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B> {
166        VecZnxBig {
167            data: self.data,
168            n: self.n,
169            cols: self.cols,
170            size: self.size,
171            _phantom: PhantomData,
172        }
173    }
174}
175
176impl<B: Backend> VecZnxBigToRef<B> for VecZnxBig<&mut [u8], B> {
177    fn to_ref(&self) -> VecZnxBig<&[u8], B> {
178        VecZnxBig {
179            data: self.data,
180            n: self.n,
181            cols: self.cols,
182            size: self.size,
183            _phantom: PhantomData,
184        }
185    }
186}
187
188impl<B: Backend> VecZnxBigToRef<B> for VecZnxBig<&[u8], B> {
189    fn to_ref(&self) -> VecZnxBig<&[u8], B> {
190        VecZnxBig {
191            data: self.data,
192            n: self.n,
193            cols: self.cols,
194            size: self.size,
195            _phantom: PhantomData,
196        }
197    }
198}
199
200impl<D: AsRef<[u8]>> fmt::Display for VecZnxBig<D, FFT64> {
201    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
202        writeln!(
203            f,
204            "VecZnxBig(n={}, cols={}, size={})",
205            self.n, self.cols, self.size
206        )?;
207
208        for col in 0..self.cols {
209            writeln!(f, "Column {}:", col)?;
210            for size in 0..self.size {
211                let coeffs = self.at(col, size);
212                write!(f, "  Size {}: [", size)?;
213
214                let max_show = 100;
215                let show_count = coeffs.len().min(max_show);
216
217                for (i, &coeff) in coeffs.iter().take(show_count).enumerate() {
218                    if i > 0 {
219                        write!(f, ", ")?;
220                    }
221                    write!(f, "{}", coeff)?;
222                }
223
224                if coeffs.len() > max_show {
225                    write!(f, ", ... ({} more)", coeffs.len() - max_show)?;
226                }
227
228                writeln!(f, "]")?;
229            }
230        }
231        Ok(())
232    }
233}