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