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