1use std::{
  2    hash::{DefaultHasher, Hasher},
  3    marker::PhantomData,
  4};
  5
  6use crate::{
  7    alloc_aligned,
  8    layouts::{Backend, Data, DataMut, DataRef, DataView, DataViewMut, DigestU64, ZnxInfos, ZnxView},
  9    oep::VmpPMatAllocBytesImpl,
 10};
 11
 12#[repr(C)]
 13#[derive(PartialEq, Eq, Hash)]
 14pub struct VmpPMat<D: Data, B: Backend> {
 15    data: D,
 16    n: usize,
 17    size: usize,
 18    rows: usize,
 19    cols_in: usize,
 20    cols_out: usize,
 21    _phantom: PhantomData<B>,
 22}
 23
 24impl<D: DataRef, B: Backend> DigestU64 for VmpPMat<D, B> {
 25    fn digest_u64(&self) -> u64 {
 26        let mut h: DefaultHasher = DefaultHasher::new();
 27        h.write(self.data.as_ref());
 28        h.write_usize(self.n);
 29        h.write_usize(self.size);
 30        h.write_usize(self.rows);
 31        h.write_usize(self.cols_in);
 32        h.write_usize(self.cols_out);
 33        h.finish()
 34    }
 35}
 36
 37impl<D: DataRef, B: Backend> ZnxView for VmpPMat<D, B> {
 38    type Scalar = B::ScalarPrep;
 39}
 40
 41impl<D: Data, B: Backend> ZnxInfos for VmpPMat<D, B> {
 42    fn cols(&self) -> usize {
 43        self.cols_in
 44    }
 45
 46    fn rows(&self) -> usize {
 47        self.rows
 48    }
 49
 50    fn n(&self) -> usize {
 51        self.n
 52    }
 53
 54    fn size(&self) -> usize {
 55        self.size
 56    }
 57
 58    fn poly_count(&self) -> usize {
 59        self.rows() * self.cols_in() * self.size() * self.cols_out()
 60    }
 61}
 62
 63impl<D: Data, B: Backend> DataView for VmpPMat<D, B> {
 64    type D = D;
 65    fn data(&self) -> &Self::D {
 66        &self.data
 67    }
 68}
 69
 70impl<D: Data, B: Backend> DataViewMut for VmpPMat<D, B> {
 71    fn data_mut(&mut self) -> &mut Self::D {
 72        &mut self.data
 73    }
 74}
 75
 76impl<D: Data, B: Backend> VmpPMat<D, B> {
 77    pub fn cols_in(&self) -> usize {
 78        self.cols_in
 79    }
 80
 81    pub fn cols_out(&self) -> usize {
 82        self.cols_out
 83    }
 84}
 85
 86impl<D: DataRef + From<Vec<u8>>, B: Backend> VmpPMat<D, B>
 87where
 88    B: VmpPMatAllocBytesImpl<B>,
 89{
 90    pub fn alloc(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> Self {
 91        let data: Vec<u8> = alloc_aligned(B::vmp_pmat_alloc_bytes_impl(
 92            n, rows, cols_in, cols_out, size,
 93        ));
 94        Self {
 95            data: data.into(),
 96            n,
 97            size,
 98            rows,
 99            cols_in,
100            cols_out,
101            _phantom: PhantomData,
102        }
103    }
104
105    pub fn from_bytes(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize, bytes: impl Into<Vec<u8>>) -> Self {
106        let data: Vec<u8> = bytes.into();
107        assert!(data.len() == B::vmp_pmat_alloc_bytes_impl(n, rows, cols_in, cols_out, size));
108        Self {
109            data: data.into(),
110            n,
111            size,
112            rows,
113            cols_in,
114            cols_out,
115            _phantom: PhantomData,
116        }
117    }
118}
119
120pub type VmpPMatOwned<B> = VmpPMat<Vec<u8>, B>;
121pub type VmpPMatRef<'a, B> = VmpPMat<&'a [u8], B>;
122
123pub trait VmpPMatToRef<B: Backend> {
124    fn to_ref(&self) -> VmpPMat<&[u8], B>;
125}
126
127impl<D: DataRef, B: Backend> VmpPMatToRef<B> for VmpPMat<D, B> {
128    fn to_ref(&self) -> VmpPMat<&[u8], B> {
129        VmpPMat {
130            data: self.data.as_ref(),
131            n: self.n,
132            rows: self.rows,
133            cols_in: self.cols_in,
134            cols_out: self.cols_out,
135            size: self.size,
136            _phantom: std::marker::PhantomData,
137        }
138    }
139}
140
141pub trait VmpPMatToMut<B: Backend> {
142    fn to_mut(&mut self) -> VmpPMat<&mut [u8], B>;
143}
144
145impl<D: DataMut, B: Backend> VmpPMatToMut<B> for VmpPMat<D, B> {
146    fn to_mut(&mut self) -> VmpPMat<&mut [u8], B> {
147        VmpPMat {
148            data: self.data.as_mut(),
149            n: self.n,
150            rows: self.rows,
151            cols_in: self.cols_in,
152            cols_out: self.cols_out,
153            size: self.size,
154            _phantom: std::marker::PhantomData,
155        }
156    }
157}
158
159impl<D: Data, B: Backend> VmpPMat<D, B> {
160    pub fn from_data(data: D, n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> Self {
161        Self {
162            data,
163            n,
164            rows,
165            cols_in,
166            cols_out,
167            size,
168            _phantom: PhantomData,
169        }
170    }
171}