1use std::marker::PhantomData;
  2
  3use rand_distr::num_traits::Zero;
  4
  5use crate::{
  6    alloc_aligned,
  7    hal::{
  8        api::{DataView, DataViewMut, ZnxInfos, ZnxView, ZnxViewMut, ZnxZero},
  9        layouts::{Backend, Data, DataMut, DataRef},
 10    },
 11};
 12
 13#[derive(PartialEq, Eq)]
 14pub struct VecZnxBig<D: Data, B: Backend> {
 15    pub(crate) data: D,
 16    pub(crate) n: usize,
 17    pub(crate) cols: usize,
 18    pub(crate) size: usize,
 19    pub(crate) max_size: usize,
 20    pub(crate) _phantom: PhantomData<B>,
 21}
 22
 23impl<D: Data, B: Backend> ZnxInfos for VecZnxBig<D, B> {
 24    fn cols(&self) -> usize {
 25        self.cols
 26    }
 27
 28    fn rows(&self) -> usize {
 29        1
 30    }
 31
 32    fn n(&self) -> usize {
 33        self.n
 34    }
 35
 36    fn size(&self) -> usize {
 37        self.size
 38    }
 39}
 40
 41impl<D: Data, B: Backend> DataView for VecZnxBig<D, B> {
 42    type D = D;
 43    fn data(&self) -> &Self::D {
 44        &self.data
 45    }
 46}
 47
 48impl<D: Data, B: Backend> DataViewMut for VecZnxBig<D, B> {
 49    fn data_mut(&mut self) -> &mut Self::D {
 50        &mut self.data
 51    }
 52}
 53
 54pub trait VecZnxBigBytesOf {
 55    fn bytes_of(n: usize, cols: usize, size: usize) -> usize;
 56}
 57
 58impl<D: DataMut, B: Backend> ZnxZero for VecZnxBig<D, B>
 59where
 60    Self: ZnxViewMut,
 61    <Self as ZnxView>::Scalar: Zero + Copy,
 62{
 63    fn zero(&mut self) {
 64        self.raw_mut().fill(<Self as ZnxView>::Scalar::zero())
 65    }
 66    fn zero_at(&mut self, i: usize, j: usize) {
 67        self.at_mut(i, j).fill(<Self as ZnxView>::Scalar::zero());
 68    }
 69}
 70
 71impl<D: DataRef + From<Vec<u8>>, B: Backend> VecZnxBig<D, B>
 72where
 73    VecZnxBig<D, B>: VecZnxBigBytesOf,
 74{
 75    pub(crate) fn new(n: usize, cols: usize, size: usize) -> Self {
 76        let data = alloc_aligned::<u8>(Self::bytes_of(n, cols, size));
 77        Self {
 78            data: data.into(),
 79            n,
 80            cols,
 81            size,
 82            max_size: size,
 83            _phantom: PhantomData,
 84        }
 85    }
 86
 87    pub(crate) fn new_from_bytes(n: usize, cols: usize, size: usize, bytes: impl Into<Vec<u8>>) -> Self {
 88        let data: Vec<u8> = bytes.into();
 89        assert!(data.len() == Self::bytes_of(n, cols, size));
 90        Self {
 91            data: data.into(),
 92            n,
 93            cols,
 94            size,
 95            max_size: size,
 96            _phantom: PhantomData,
 97        }
 98    }
 99}
100
101impl<D: Data, B: Backend> VecZnxBig<D, B> {
102    pub(crate) fn from_data(data: D, n: usize, cols: usize, size: usize) -> Self {
103        Self {
104            data,
105            n,
106            cols,
107            size,
108            max_size: size,
109            _phantom: PhantomData,
110        }
111    }
112}
113
114pub type VecZnxBigOwned<B> = VecZnxBig<Vec<u8>, B>;
115
116pub trait VecZnxBigToRef<B: Backend> {
117    fn to_ref(&self) -> VecZnxBig<&[u8], B>;
118}
119
120impl<D: DataRef, B: Backend> VecZnxBigToRef<B> for VecZnxBig<D, B> {
121    fn to_ref(&self) -> VecZnxBig<&[u8], B> {
122        VecZnxBig {
123            data: self.data.as_ref(),
124            n: self.n,
125            cols: self.cols,
126            size: self.size,
127            max_size: self.max_size,
128            _phantom: std::marker::PhantomData,
129        }
130    }
131}
132
133pub trait VecZnxBigToMut<B: Backend> {
134    fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B>;
135}
136
137impl<D: DataMut, B: Backend> VecZnxBigToMut<B> for VecZnxBig<D, B> {
138    fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B> {
139        VecZnxBig {
140            data: self.data.as_mut(),
141            n: self.n,
142            cols: self.cols,
143            size: self.size,
144            max_size: self.max_size,
145            _phantom: std::marker::PhantomData,
146        }
147    }
148}