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};
  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_apply_dft_tmp_bytes(
114        &self,
115        c_limbs: usize,
116        a_limbs: usize,
117        rows: usize,
118        cols: usize,
119    ) -> usize {
120        unsafe {
121            vmp_apply_dft_tmp_bytes(
122                self.0,
123                c_limbs as u64,
124                a_limbs as u64,
125                rows as u64,
126                cols as u64,
127            ) as usize
128        }
129    }
130
131    pub fn vmp_apply_dft(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, buf: &mut [u8]) {
132        unsafe {
133            vmp_apply_dft(
134                self.0,
135                c.0,
136                c.limbs() as u64,
137                a.as_ptr(),
138                a.limbs() as u64,
139                a.n() as u64,
140                b.data(),
141                b.rows() as u64,
142                b.cols() as u64,
143                buf.as_mut_ptr(),
144            )
145        }
146    }
147
148    pub fn vmp_apply_dft_to_dft_tmp_bytes(
149        &self,
150        c_limbs: usize,
151        a_limbs: usize,
152        rows: usize,
153        cols: usize,
154    ) -> usize {
155        unsafe {
156            vmp_apply_dft_to_dft_tmp_bytes(
157                self.0,
158                c_limbs as u64,
159                a_limbs as u64,
160                rows as u64,
161                cols as u64,
162            ) as usize
163        }
164    }
165
166    pub fn vmp_apply_dft_to_dft(
167        &self,
168        c: &mut VecZnxDft,
169        a: &VecZnxDft,
170        b: &VmpPMat,
171        buf: &mut [u8],
172    ) {
173        unsafe {
174            vmp_apply_dft_to_dft(
175                self.0,
176                c.0,
177                c.limbs() as u64,
178                a.0,
179                a.limbs() as u64,
180                b.data(),
181                b.rows() as u64,
182                b.cols() as u64,
183                buf.as_mut_ptr(),
184            )
185        }
186    }
187
188    pub fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, buf: &mut [u8]) {
189        unsafe {
190            vmp_apply_dft_to_dft(
191                self.0,
192                b.0,
193                b.limbs() as u64,
194                b.0,
195                b.limbs() as u64,
196                a.data(),
197                a.rows() as u64,
198                a.cols() as u64,
199                buf.as_mut_ptr(),
200            )
201        }
202    }
203}
204
205pub struct Matrix3D<T> {
206    pub data: Vec<T>,
207    pub rows: usize,
208    pub cols: usize,
209    pub n: usize,
210}
211
212impl<T: Default + Clone> Matrix3D<T> {
213    pub fn new(rows: usize, cols: usize, n: usize) -> Self {
214        let size = rows * cols * n;
215        Self {
216            data: vec![T::default(); size],
217            rows,
218            cols,
219            n,
220        }
221    }
222
223    pub fn at(&self, row: usize, col: usize) -> &[T] {
224        assert!(row <= self.rows && col <= self.cols);
225        let idx: usize = col * (self.n * self.rows) + row * self.n;
226        &self.data[idx..idx + self.n]
227    }
228
229    pub fn at_mut(&mut self, row: usize, col: usize) -> &mut [T] {
230        assert!(row <= self.rows && col <= self.cols);
231        let idx: usize = col * (self.n * self.rows) + row * self.n;
232        &mut self.data[idx..idx + self.n]
233    }
234}