1use crate::{
  2    api::{CnvPVecBytesOf, ModuleN, SvpPPolBytesOf, VecZnxBigBytesOf, VecZnxDftBytesOf, VmpPMatBytesOf},
  3    layouts::{Backend, CnvPVecL, CnvPVecR, MatZnx, ScalarZnx, Scratch, SvpPPol, VecZnx, VecZnxBig, VecZnxDft, VmpPMat},
  4};
  5
  6/// Allocates a new [crate::layouts::ScratchOwned] of `size` aligned bytes.
  7pub trait ScratchOwnedAlloc<B: Backend> {
  8    fn alloc(size: usize) -> Self;
  9}
 10
 11/// Borrows a slice of bytes into a [Scratch].
 12pub trait ScratchOwnedBorrow<B: Backend> {
 13    fn borrow(&mut self) -> &mut Scratch<B>;
 14}
 15
 16/// Wrap an array of mutable borrowed bytes into a [Scratch].
 17pub trait ScratchFromBytes<B: Backend> {
 18    fn from_bytes(data: &mut [u8]) -> &mut Scratch<B>;
 19}
 20
 21/// Returns how many bytes left can be taken from the scratch.
 22pub trait ScratchAvailable {
 23    fn available(&self) -> usize;
 24}
 25
 26/// Takes a slice of bytes from a [Scratch] and return a new [Scratch] minus the taken array of bytes.
 27pub trait TakeSlice {
 28    fn take_slice<T>(&mut self, len: usize) -> (&mut [T], &mut Self);
 29}
 30
 31impl<BE: Backend> Scratch<BE>
 32where
 33    Self: TakeSlice + ScratchAvailable + ScratchFromBytes<BE>,
 34{
 35    pub fn split_at_mut(&mut self, len: usize) -> (&mut Scratch<BE>, &mut Self) {
 36        let (take_slice, rem_slice) = self.take_slice(len);
 37        (Self::from_bytes(take_slice), rem_slice)
 38    }
 39
 40    pub fn split_mut(&mut self, n: usize, len: usize) -> (Vec<&mut Scratch<BE>>, &mut Self) {
 41        assert!(self.available() >= n * len);
 42        let mut scratches: Vec<&mut Scratch<BE>> = Vec::with_capacity(n);
 43        let mut scratch: &mut Scratch<BE> = self;
 44        for _ in 0..n {
 45            let (tmp, scratch_new) = scratch.split_at_mut(len);
 46            scratch = scratch_new;
 47            scratches.push(tmp);
 48        }
 49        (scratches, scratch)
 50    }
 51}
 52
 53impl<B: Backend> ScratchTakeBasic for Scratch<B> where Self: TakeSlice + ScratchFromBytes<B> {}
 54
 55pub trait ScratchTakeBasic
 56where
 57    Self: TakeSlice,
 58{
 59    fn take_cnv_pvec_left<M, B: Backend>(&mut self, module: &M, cols: usize, size: usize) -> (CnvPVecL<&mut [u8], B>, &mut Self)
 60    where
 61        M: ModuleN + CnvPVecBytesOf,
 62    {
 63        let (take_slice, rem_slice) = self.take_slice(module.bytes_of_cnv_pvec_left(cols, size));
 64        (CnvPVecL::from_data(take_slice, module.n(), cols, size), rem_slice)
 65    }
 66
 67    fn take_cnv_pvec_right<M, B: Backend>(&mut self, module: &M, cols: usize, size: usize) -> (CnvPVecR<&mut [u8], B>, &mut Self)
 68    where
 69        M: ModuleN + CnvPVecBytesOf,
 70    {
 71        let (take_slice, rem_slice) = self.take_slice(module.bytes_of_cnv_pvec_right(cols, size));
 72        (CnvPVecR::from_data(take_slice, module.n(), cols, size), rem_slice)
 73    }
 74
 75    fn take_scalar_znx(&mut self, n: usize, cols: usize) -> (ScalarZnx<&mut [u8]>, &mut Self) {
 76        let (take_slice, rem_slice) = self.take_slice(ScalarZnx::bytes_of(n, cols));
 77        (ScalarZnx::from_data(take_slice, n, cols), rem_slice)
 78    }
 79
 80    fn take_svp_ppol<M, B: Backend>(&mut self, module: &M, cols: usize) -> (SvpPPol<&mut [u8], B>, &mut Self)
 81    where
 82        M: SvpPPolBytesOf + ModuleN,
 83    {
 84        let (take_slice, rem_slice) = self.take_slice(module.bytes_of_svp_ppol(cols));
 85        (SvpPPol::from_data(take_slice, module.n(), cols), rem_slice)
 86    }
 87
 88    fn take_vec_znx(&mut self, n: usize, cols: usize, size: usize) -> (VecZnx<&mut [u8]>, &mut Self) {
 89        let (take_slice, rem_slice) = self.take_slice(VecZnx::bytes_of(n, cols, size));
 90        (VecZnx::from_data(take_slice, n, cols, size), rem_slice)
 91    }
 92
 93    fn take_vec_znx_big<M, B: Backend>(&mut self, module: &M, cols: usize, size: usize) -> (VecZnxBig<&mut [u8], B>, &mut Self)
 94    where
 95        M: VecZnxBigBytesOf + ModuleN,
 96    {
 97        let (take_slice, rem_slice) = self.take_slice(module.bytes_of_vec_znx_big(cols, size));
 98        (VecZnxBig::from_data(take_slice, module.n(), cols, size), rem_slice)
 99    }
100
101    fn take_vec_znx_dft<M, B: Backend>(&mut self, module: &M, cols: usize, size: usize) -> (VecZnxDft<&mut [u8], B>, &mut Self)
102    where
103        M: VecZnxDftBytesOf + ModuleN,
104    {
105        let (take_slice, rem_slice) = self.take_slice(module.bytes_of_vec_znx_dft(cols, size));
106
107        (VecZnxDft::from_data(take_slice, module.n(), cols, size), rem_slice)
108    }
109
110    fn take_vec_znx_dft_slice<M, B: Backend>(
111        &mut self,
112        module: &M,
113        len: usize,
114        cols: usize,
115        size: usize,
116    ) -> (Vec<VecZnxDft<&mut [u8], B>>, &mut Self)
117    where
118        M: VecZnxDftBytesOf + ModuleN,
119    {
120        let mut scratch: &mut Self = self;
121        let mut slice: Vec<VecZnxDft<&mut [u8], B>> = Vec::with_capacity(len);
122        for _ in 0..len {
123            let (znx, new_scratch) = scratch.take_vec_znx_dft(module, cols, size);
124            scratch = new_scratch;
125            slice.push(znx);
126        }
127        (slice, scratch)
128    }
129
130    fn take_vec_znx_slice(&mut self, len: usize, n: usize, cols: usize, size: usize) -> (Vec<VecZnx<&mut [u8]>>, &mut Self) {
131        let mut scratch: &mut Self = self;
132        let mut slice: Vec<VecZnx<&mut [u8]>> = Vec::with_capacity(len);
133        for _ in 0..len {
134            let (znx, new_scratch) = scratch.take_vec_znx(n, cols, size);
135            scratch = new_scratch;
136            slice.push(znx);
137        }
138        (slice, scratch)
139    }
140
141    fn take_vmp_pmat<M, B: Backend>(
142        &mut self,
143        module: &M,
144        rows: usize,
145        cols_in: usize,
146        cols_out: usize,
147        size: usize,
148    ) -> (VmpPMat<&mut [u8], B>, &mut Self)
149    where
150        M: VmpPMatBytesOf + ModuleN,
151    {
152        let (take_slice, rem_slice) = self.take_slice(module.bytes_of_vmp_pmat(rows, cols_in, cols_out, size));
153        (
154            VmpPMat::from_data(take_slice, module.n(), rows, cols_in, cols_out, size),
155            rem_slice,
156        )
157    }
158
159    fn take_mat_znx(
160        &mut self,
161        n: usize,
162        rows: usize,
163        cols_in: usize,
164        cols_out: usize,
165        size: usize,
166    ) -> (MatZnx<&mut [u8]>, &mut Self) {
167        let (take_slice, rem_slice) = self.take_slice(MatZnx::bytes_of(n, rows, cols_in, cols_out, size));
168        (MatZnx::from_data(take_slice, n, rows, cols_in, cols_out, size), rem_slice)
169    }
170}