1use crate::bindings::{
  2    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(pub *mut vmp_pmat_t, pub usize, pub usize);
  9
 10impl VmpPMat {
 11    pub fn rows(&self) -> usize {
 12        self.1
 13    }
 14
 15    pub fn cols(&self) -> usize {
 16        self.2
 17    }
 18}
 19
 20impl Module {
 21    pub fn new_vmp_pmat(&self, rows: usize, cols: usize) -> VmpPMat {
 22        unsafe { VmpPMat(new_vmp_pmat(self.0, rows as u64, cols as u64), rows, cols) }
 23    }
 24
 25    pub fn vmp_prepare_contiguous_tmp_bytes(&self, rows: usize, cols: usize) -> usize {
 26        unsafe { vmp_prepare_contiguous_tmp_bytes(self.0, rows as u64, cols as u64) as usize }
 27    }
 28
 29    pub fn vmp_prepare_contiguous(&self, b: &mut VmpPMat, a: &[i64], buf: &mut [u8]) {
 30        unsafe {
 31            vmp_prepare_contiguous(
 32                self.0,
 33                b.0,
 34                a.as_ptr(),
 35                b.1 as u64,
 36                b.2 as u64,
 37                buf.as_mut_ptr(),
 38            );
 39        }
 40    }
 41
 42    pub fn vmp_apply_dft_tmp_bytes(
 43        &self,
 44        c_limbs: usize,
 45        a_limbs: usize,
 46        rows: usize,
 47        cols: usize,
 48    ) -> usize {
 49        unsafe {
 50            vmp_apply_dft_tmp_bytes(
 51                self.0,
 52                c_limbs as u64,
 53                a_limbs as u64,
 54                rows as u64,
 55                cols as u64,
 56            ) as usize
 57        }
 58    }
 59
 60    pub fn vmp_apply_dft(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, buf: &mut [u8]) {
 61        unsafe {
 62            vmp_apply_dft(
 63                self.0,
 64                c.0,
 65                c.limbs() as u64,
 66                a.as_ptr(),
 67                a.limbs() as u64,
 68                a.n() as u64,
 69                b.0,
 70                b.rows() as u64,
 71                b.cols() as u64,
 72                buf.as_mut_ptr(),
 73            )
 74        }
 75    }
 76
 77    pub fn vmp_apply_dft_to_dft_tmp_bytes(
 78        &self,
 79        c_limbs: usize,
 80        a_limbs: usize,
 81        rows: usize,
 82        cols: usize,
 83    ) -> usize {
 84        unsafe {
 85            vmp_apply_dft_to_dft_tmp_bytes(
 86                self.0,
 87                c_limbs as u64,
 88                a_limbs as u64,
 89                rows as u64,
 90                cols as u64,
 91            ) as usize
 92        }
 93    }
 94
 95    pub fn vmp_apply_dft_to_dft(
 96        &self,
 97        c: &mut VecZnxDft,
 98        a: &VecZnxDft,
 99        b: &VmpPMat,
100        buf: &mut [u8],
101    ) {
102        unsafe {
103            vmp_apply_dft_to_dft(
104                self.0,
105                c.0,
106                c.limbs() as u64,
107                a.0,
108                a.limbs() as u64,
109                b.0,
110                b.rows() as u64,
111                b.cols() as u64,
112                buf.as_mut_ptr(),
113            )
114        }
115    }
116
117    pub fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, buf: &mut [u8]) {
118        unsafe {
119            vmp_apply_dft_to_dft(
120                self.0,
121                b.0,
122                b.limbs() as u64,
123                b.0,
124                b.limbs() as u64,
125                a.0,
126                a.rows() as u64,
127                a.cols() as u64,
128                buf.as_mut_ptr(),
129            )
130        }
131    }
132}
133
134pub struct Matrix3D<T> {
135    pub data: Vec<T>,
136    pub rows: usize,
137    pub cols: usize,
138    pub n: usize,
139}
140
141impl<T: Default + Clone> Matrix3D<T> {
142    pub fn new(rows: usize, cols: usize, n: usize) -> Self {
143        let size = rows * cols * n;
144        Self {
145            data: vec![T::default(); size],
146            rows,
147            cols,
148            n,
149        }
150    }
151
152    pub fn at(&self, row: usize, col: usize) -> &[T] {
153        assert!(row <= self.rows && col <= self.cols);
154        let idx: usize = col * (self.n * self.rows) + row * self.n;
155        &self.data[idx..idx + self.n]
156    }
157
158    pub fn at_mut(&mut self, row: usize, col: usize) -> &mut [T] {
159        assert!(row <= self.rows && col <= self.cols);
160        let idx: usize = col * (self.n * self.rows) + row * self.n;
161        &mut self.data[idx..idx + self.n]
162    }
163}