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: AsMut<[u8]> + AsRef<[u8]>> 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: &C, a_col: usize)
114    where
115        C: 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
147impl<D, B: Backend> VecZnxBigToRef<B> for VecZnxBig<D, B>
148where
149    D: AsRef<[u8]>,
150    B: Backend,
151{
152    fn to_ref(&self) -> VecZnxBig<&[u8], B> {
153        VecZnxBig {
154            data: self.data.as_ref(),
155            n: self.n,
156            cols: self.cols,
157            size: self.size,
158            _phantom: std::marker::PhantomData,
159        }
160    }
161}
162
163pub trait VecZnxBigToMut<B: Backend> {
164    fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B>;
165}
166
167impl<D, B: Backend> VecZnxBigToMut<B> for VecZnxBig<D, B>
168where
169    D: AsRef<[u8]> + AsMut<[u8]>,
170    B: Backend,
171{
172    fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B> {
173        VecZnxBig {
174            data: self.data.as_mut(),
175            n: self.n,
176            cols: self.cols,
177            size: self.size,
178            _phantom: std::marker::PhantomData,
179        }
180    }
181}
182
183impl<D: AsRef<[u8]>> fmt::Display for VecZnxBig<D, FFT64> {
184    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
185        writeln!(
186            f,
187            "VecZnxBig(n={}, cols={}, size={})",
188            self.n, self.cols, self.size
189        )?;
190
191        for col in 0..self.cols {
192            writeln!(f, "Column {}:", col)?;
193            for size in 0..self.size {
194                let coeffs = self.at(col, size);
195                write!(f, "  Size {}: [", size)?;
196
197                let max_show = 100;
198                let show_count = coeffs.len().min(max_show);
199
200                for (i, &coeff) in coeffs.iter().take(show_count).enumerate() {
201                    if i > 0 {
202                        write!(f, ", ")?;
203                    }
204                    write!(f, "{}", coeff)?;
205                }
206
207                if coeffs.len() > max_show {
208                    write!(f, ", ... ({} more)", coeffs.len() - max_show)?;
209                }
210
211                writeln!(f, "]")?;
212            }
213        }
214        Ok(())
215    }
216}