1use crate::ffi::vmp::{
  2    delete_vmp_pmat, new_vmp_pmat, vmp_apply_dft, vmp_apply_dft_tmp_bytes, vmp_apply_dft_to_dft,
  3    vmp_apply_dft_to_dft_tmp_bytes, vmp_pmat_t, vmp_prepare_contiguous,
  4    vmp_prepare_contiguous_tmp_bytes, vmp_prepare_dblptr,
  5};
  6use crate::{Module, VecZnx, VecZnxDft};
  7
  8pub struct VmpPMat {
  9    pub data: *mut vmp_pmat_t,
 10    pub rows: usize,
 11    pub cols: usize,
 12    pub n: usize,
 13}
 14
 15impl VmpPMat {
 16    pub fn data(&self) -> *mut vmp_pmat_t {
 17        self.data
 18    }
 19
 20    pub fn rows(&self) -> usize {
 21        self.rows
 22    }
 23
 24    pub fn cols(&self) -> usize {
 25        self.cols
 26    }
 27
 28    pub fn n(&self) -> usize {
 29        self.n
 30    }
 31
 32    pub fn as_f64(&self) -> &[f64] {
 33        let ptr: *const f64 = self.data as *const f64;
 34        let len: usize = (self.rows() * self.cols() * self.n() * 8) / std::mem::size_of::<f64>();
 35        unsafe { &std::slice::from_raw_parts(ptr, len) }
 36    }
 37
 38    pub fn get_addr(&self, row: usize, col: usize, blk: usize) -> &[f64] {
 39        let nrows: usize = self.rows();
 40        let ncols: usize = self.cols();
 41        if col == (ncols - 1) && (ncols & 1 == 1) {
 42            &self.as_f64()[blk * nrows * ncols * 8 + col * nrows * 8 + row * 8..]
 43        } else {
 44            &self.as_f64()[blk * nrows * ncols * 8
 45                + (col / 2) * (2 * nrows) * 8
 46                + row * 2 * 8
 47                + (col % 2) * 8..]
 48        }
 49    }
 50
 51    pub fn at(&self, row: usize, col: usize) -> Vec<f64> {
 52        //assert!(row <= self.rows && col <= self.cols);
 53
 54        let mut res: Vec<f64> = vec![f64::default(); self.n];
 55
 56        if self.n < 8 {
 57            res.copy_from_slice(
 58                &self.as_f64()[(row + col * self.rows()) * self.n()
 59                    ..(row + col * self.rows()) * (self.n() + 1)],
 60            );
 61        } else {
 62            (0..self.n >> 3).for_each(|blk| {
 63                res[blk * 8..(blk + 1) * 8].copy_from_slice(&self.get_addr(row, col, blk)[..8]);
 64            });
 65        }
 66
 67        res
 68    }
 69
 70    pub fn at_mut(&self, row: usize, col: usize) -> &mut [f64] {
 71        assert!(row <= self.rows && col <= self.cols);
 72        let idx: usize = col * (self.n / 2 * self.rows) + row * (self.n >> 1);
 73        let ptr: *mut f64 = self.data as *mut f64;
 74        let len: usize = (self.rows() * self.cols() * self.n() * 8) / std::mem::size_of::<f64>();
 75        unsafe { &mut std::slice::from_raw_parts_mut(ptr, len)[idx..idx + self.n] }
 76    }
 77
 78    pub fn delete(self) {
 79        unsafe { delete_vmp_pmat(self.data) };
 80        drop(self);
 81    }
 82}
 83
 84impl Module {
 85    pub fn new_vmp_pmat(&self, rows: usize, cols: usize) -> VmpPMat {
 86        unsafe {
 87            VmpPMat {
 88                data: new_vmp_pmat(self.0, rows as u64, cols as u64),
 89                rows,
 90                cols,
 91                n: self.n(),
 92            }
 93        }
 94    }
 95
 96    pub fn vmp_prepare_contiguous_tmp_bytes(&self, rows: usize, cols: usize) -> usize {
 97        unsafe { vmp_prepare_contiguous_tmp_bytes(self.0, rows as u64, cols as u64) as usize }
 98    }
 99
100    pub fn vmp_prepare_contiguous(&self, b: &mut VmpPMat, a: &[i64], buf: &mut [u8]) {
101        unsafe {
102            vmp_prepare_contiguous(
103                self.0,
104                b.data(),
105                a.as_ptr(),
106                b.rows() as u64,
107                b.cols() as u64,
108                buf.as_mut_ptr(),
109            );
110        }
111    }
112
113    pub fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &Vec<&Vec<i64>>, buf: &mut [u8]) {
114        let ptrs: Vec<*const i64> = a.iter().map(|v| v.as_ptr()).collect();
115
116        unsafe {
117            vmp_prepare_dblptr(
118                self.0,
119                b.data(),
120                ptrs.as_ptr(),
121                b.rows() as u64,
122                b.cols() as u64,
123                buf.as_mut_ptr(),
124            );
125        }
126    }
127
128    pub fn vmp_apply_dft_tmp_bytes(
129        &self,
130        c_limbs: usize,
131        a_limbs: usize,
132        rows: usize,
133        cols: usize,
134    ) -> usize {
135        unsafe {
136            vmp_apply_dft_tmp_bytes(
137                self.0,
138                c_limbs as u64,
139                a_limbs as u64,
140                rows as u64,
141                cols as u64,
142            ) as usize
143        }
144    }
145
146    pub fn vmp_apply_dft(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, buf: &mut [u8]) {
147        unsafe {
148            vmp_apply_dft(
149                self.0,
150                c.0,
151                c.limbs() as u64,
152                a.as_ptr(),
153                a.limbs() as u64,
154                a.n() as u64,
155                b.data(),
156                b.rows() as u64,
157                b.cols() as u64,
158                buf.as_mut_ptr(),
159            )
160        }
161    }
162
163    pub fn vmp_apply_dft_to_dft_tmp_bytes(
164        &self,
165        c_limbs: usize,
166        a_limbs: usize,
167        rows: usize,
168        cols: usize,
169    ) -> usize {
170        unsafe {
171            vmp_apply_dft_to_dft_tmp_bytes(
172                self.0,
173                c_limbs as u64,
174                a_limbs as u64,
175                rows as u64,
176                cols as u64,
177            ) as usize
178        }
179    }
180
181    pub fn vmp_apply_dft_to_dft(
182        &self,
183        c: &mut VecZnxDft,
184        a: &VecZnxDft,
185        b: &VmpPMat,
186        buf: &mut [u8],
187    ) {
188        unsafe {
189            vmp_apply_dft_to_dft(
190                self.0,
191                c.0,
192                c.limbs() as u64,
193                a.0,
194                a.limbs() as u64,
195                b.data(),
196                b.rows() as u64,
197                b.cols() as u64,
198                buf.as_mut_ptr(),
199            )
200        }
201    }
202
203    pub fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, buf: &mut [u8]) {
204        unsafe {
205            vmp_apply_dft_to_dft(
206                self.0,
207                b.0,
208                b.limbs() as u64,
209                b.0,
210                b.limbs() as u64,
211                a.data(),
212                a.rows() as u64,
213                a.cols() as u64,
214                buf.as_mut_ptr(),
215            )
216        }
217    }
218}
219
220pub struct Matrix3D<T> {
221    pub data: Vec<T>,
222    pub rows: usize,
223    pub cols: usize,
224    pub n: usize,
225}
226
227impl<T: Default + Clone + std::marker::Copy> Matrix3D<T> {
228    pub fn new(rows: usize, cols: usize, n: usize) -> Self {
229        let size = rows * cols * n;
230        Self {
231            data: vec![T::default(); size],
232            rows,
233            cols,
234            n,
235        }
236    }
237
238    pub fn at(&self, row: usize, col: usize) -> &[T] {
239        assert!(row <= self.rows && col <= self.cols);
240        let idx: usize = col * (self.n * self.rows) + row * self.n;
241        &self.data[idx..idx + self.n]
242    }
243
244    pub fn at_mut(&mut self, row: usize, col: usize) -> &mut [T] {
245        assert!(row <= self.rows && col <= self.cols);
246        let idx: usize = col * (self.n * self.rows) + row * self.n;
247        &mut self.data[idx..idx + self.n]
248    }
249
250    pub fn set_col(&mut self, col: usize, a: &[T]) {
251        let idx: usize = col * (self.n * self.rows);
252        self.data[idx..idx + self.rows * self.n].copy_from_slice(a);
253    }
254}