1use crate::ffi::vec_znx_big;
  2use crate::znx_base::{ZnxInfos, ZnxView};
  3use crate::{Backend, DataView, DataViewMut, FFT64, Module, VecZnx, 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    // Consumes the VecZnxBig to return a VecZnx.
102    // Useful when no normalization is needed.
103    pub fn to_vec_znx_small(self) -> VecZnx<D> {
104        VecZnx {
105            data: self.data,
106            n: self.n,
107            cols: self.cols,
108            size: self.size,
109        }
110    }
111
112    /// Extracts the a_col-th column of 'a' and stores it on the self_col-th column [Self].
113    pub fn extract_column<C>(&mut self, self_col: usize, a: &VecZnxBig<C, FFT64>, a_col: usize)
114    where
115        VecZnxBig<C, FFT64>: VecZnxBigToRef<FFT64> + ZnxInfos,
116    {
117        #[cfg(debug_assertions)]
118        {
119            assert!(self_col < self.cols());
120            assert!(a_col < a.cols());
121        }
122
123        let min_size: usize = self.size.min(a.size());
124        let max_size: usize = self.size;
125
126        let mut self_mut: VecZnxBig<&mut [u8], FFT64> = self.to_mut();
127        let a_ref: VecZnxBig<&[u8], FFT64> = a.to_ref();
128
129        (0..min_size).for_each(|i: usize| {
130            self_mut
131                .at_mut(self_col, i)
132                .copy_from_slice(a_ref.at(a_col, i));
133        });
134
135        (min_size..max_size).for_each(|i| {
136            self_mut.zero_at(self_col, i);
137        });
138    }
139}
140
141pub type VecZnxBigOwned<B> = VecZnxBig<Vec<u8>, B>;
142
143pub trait VecZnxBigToRef<B: Backend> {
144    fn to_ref(&self) -> VecZnxBig<&[u8], B>;
145}
146
147pub trait VecZnxBigToMut<B: Backend> {
148    fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B>;
149}
150
151impl<B: Backend> VecZnxBigToMut<B> for VecZnxBig<Vec<u8>, B> {
152    fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B> {
153        VecZnxBig {
154            data: self.data.as_mut_slice(),
155            n: self.n,
156            cols: self.cols,
157            size: self.size,
158            _phantom: PhantomData,
159        }
160    }
161}
162
163impl<B: Backend> VecZnxBigToRef<B> for VecZnxBig<Vec<u8>, B> {
164    fn to_ref(&self) -> VecZnxBig<&[u8], B> {
165        VecZnxBig {
166            data: self.data.as_slice(),
167            n: self.n,
168            cols: self.cols,
169            size: self.size,
170            _phantom: PhantomData,
171        }
172    }
173}
174
175impl<B: Backend> VecZnxBigToMut<B> for VecZnxBig<&mut [u8], B> {
176    fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B> {
177        VecZnxBig {
178            data: self.data,
179            n: self.n,
180            cols: self.cols,
181            size: self.size,
182            _phantom: PhantomData,
183        }
184    }
185}
186
187impl<B: Backend> VecZnxBigToRef<B> for VecZnxBig<&mut [u8], B> {
188    fn to_ref(&self) -> VecZnxBig<&[u8], B> {
189        VecZnxBig {
190            data: self.data,
191            n: self.n,
192            cols: self.cols,
193            size: self.size,
194            _phantom: PhantomData,
195        }
196    }
197}
198
199impl<B: Backend> VecZnxBigToRef<B> for VecZnxBig<&[u8], B> {
200    fn to_ref(&self) -> VecZnxBig<&[u8], B> {
201        VecZnxBig {
202            data: self.data,
203            n: self.n,
204            cols: self.cols,
205            size: self.size,
206            _phantom: PhantomData,
207        }
208    }
209}
210
211impl<D: AsRef<[u8]>> fmt::Display for VecZnxBig<D, FFT64> {
212    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
213        writeln!(
214            f,
215            "VecZnxBig(n={}, cols={}, size={})",
216            self.n, self.cols, self.size
217        )?;
218
219        for col in 0..self.cols {
220            writeln!(f, "Column {}:", col)?;
221            for size in 0..self.size {
222                let coeffs = self.at(col, size);
223                write!(f, "  Size {}: [", size)?;
224
225                let max_show = 100;
226                let show_count = coeffs.len().min(max_show);
227
228                for (i, &coeff) in coeffs.iter().take(show_count).enumerate() {
229                    if i > 0 {
230                        write!(f, ", ")?;
231                    }
232                    write!(f, "{}", coeff)?;
233                }
234
235                if coeffs.len() > max_show {
236                    write!(f, ", ... ({} more)", coeffs.len() - max_show)?;
237                }
238
239                writeln!(f, "]")?;
240            }
241        }
242        Ok(())
243    }
244}