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_bytes_of_impl(n, rows, cols_in, cols_out, size));
 92        Self {
 93            data: data.into(),
 94            n,
 95            size,
 96            rows,
 97            cols_in,
 98            cols_out,
 99            _phantom: PhantomData,
100        }
101    }
102
103    pub fn from_bytes(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize, bytes: impl Into<Vec<u8>>) -> Self {
104        let data: Vec<u8> = bytes.into();
105        assert!(data.len() == B::vmp_pmat_bytes_of_impl(n, rows, cols_in, cols_out, size));
106        Self {
107            data: data.into(),
108            n,
109            size,
110            rows,
111            cols_in,
112            cols_out,
113            _phantom: PhantomData,
114        }
115    }
116}
117
118pub type VmpPMatOwned<B> = VmpPMat<Vec<u8>, B>;
119pub type VmpPMatRef<'a, B> = VmpPMat<&'a [u8], B>;
120
121pub trait VmpPMatToRef<B: Backend> {
122    fn to_ref(&self) -> VmpPMat<&[u8], B>;
123}
124
125impl<D: DataRef, B: Backend> VmpPMatToRef<B> for VmpPMat<D, B> {
126    fn to_ref(&self) -> VmpPMat<&[u8], B> {
127        VmpPMat {
128            data: self.data.as_ref(),
129            n: self.n,
130            rows: self.rows,
131            cols_in: self.cols_in,
132            cols_out: self.cols_out,
133            size: self.size,
134            _phantom: std::marker::PhantomData,
135        }
136    }
137}
138
139pub trait VmpPMatToMut<B: Backend> {
140    fn to_mut(&mut self) -> VmpPMat<&mut [u8], B>;
141}
142
143impl<D: DataMut, B: Backend> VmpPMatToMut<B> for VmpPMat<D, B> {
144    fn to_mut(&mut self) -> VmpPMat<&mut [u8], B> {
145        VmpPMat {
146            data: self.data.as_mut(),
147            n: self.n,
148            rows: self.rows,
149            cols_in: self.cols_in,
150            cols_out: self.cols_out,
151            size: self.size,
152            _phantom: std::marker::PhantomData,
153        }
154    }
155}
156
157impl<D: Data, B: Backend> VmpPMat<D, B> {
158    pub fn from_data(data: D, n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> Self {
159        Self {
160            data,
161            n,
162            rows,
163            cols_in,
164            cols_out,
165            size,
166            _phantom: PhantomData,
167        }
168    }
169}