1use crate::{
  2    api::{ModuleN, SvpPPolBytesOf, VecZnxBigBytesOf, VecZnxDftBytesOf, VmpPMatBytesOf},
  3    layouts::{Backend, 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<B: Backend> ScratchTakeBasic for Scratch<B> where Self: TakeSlice {}
 32
 33pub trait ScratchTakeBasic
 34where
 35    Self: TakeSlice,
 36{
 37    fn take_scalar_znx(&mut self, n: usize, cols: usize) -> (ScalarZnx<&mut [u8]>, &mut Self)
 38    {
 39        let (take_slice, rem_slice) = self.take_slice(ScalarZnx::bytes_of(n, cols));
 40        (
 41            ScalarZnx::from_data(take_slice, n, cols),
 42            rem_slice,
 43        )
 44    }
 45
 46    fn take_svp_ppol<M, B: Backend>(&mut self, module: &M, cols: usize) -> (SvpPPol<&mut [u8], B>, &mut Self)
 47    where
 48        M: SvpPPolBytesOf + ModuleN,
 49    {
 50        let (take_slice, rem_slice) = self.take_slice(module.bytes_of_svp_ppol(cols));
 51        (SvpPPol::from_data(take_slice, module.n(), cols), rem_slice)
 52    }
 53
 54    fn take_vec_znx(&mut self, n: usize, cols: usize, size: usize) -> (VecZnx<&mut [u8]>, &mut Self){
 55        let (take_slice, rem_slice) = self.take_slice(VecZnx::bytes_of(n, cols, size));
 56        (
 57            VecZnx::from_data(take_slice, n, cols, size),
 58            rem_slice,
 59        )
 60    }
 61
 62    fn take_vec_znx_big<M, B: Backend>(&mut self, module: &M, cols: usize, size: usize) -> (VecZnxBig<&mut [u8], B>, &mut Self)
 63    where
 64        M: VecZnxBigBytesOf + ModuleN,
 65    {
 66        let (take_slice, rem_slice) = self.take_slice(module.bytes_of_vec_znx_big(cols, size));
 67        (
 68            VecZnxBig::from_data(take_slice, module.n(), cols, size),
 69            rem_slice,
 70        )
 71    }
 72
 73    fn take_vec_znx_dft<M, B: Backend>(&mut self, module: &M, cols: usize, size: usize) -> (VecZnxDft<&mut [u8], B>, &mut Self)
 74    where
 75        M: VecZnxDftBytesOf + ModuleN,
 76    {
 77        let (take_slice, rem_slice) = self.take_slice(module.bytes_of_vec_znx_dft(cols, size));
 78
 79        (
 80            VecZnxDft::from_data(take_slice, module.n(), cols, size),
 81            rem_slice,
 82        )
 83    }
 84
 85    fn take_vec_znx_dft_slice<M, B: Backend>(
 86        &mut self,
 87        module: &M,
 88        len: usize,
 89        cols: usize,
 90        size: usize,
 91    ) -> (Vec<VecZnxDft<&mut [u8], B>>, &mut Self)
 92    where
 93        M: VecZnxDftBytesOf + ModuleN,
 94    {
 95        let mut scratch: &mut Self = self;
 96        let mut slice: Vec<VecZnxDft<&mut [u8], B>> = Vec::with_capacity(len);
 97        for _ in 0..len {
 98            let (znx, new_scratch) = scratch.take_vec_znx_dft(module, cols, size);
 99            scratch = new_scratch;
100            slice.push(znx);
101        }
102        (slice, scratch)
103    }
104
105    fn take_vec_znx_slice(&mut self, n: usize, len: usize, cols: usize, size: usize) -> (Vec<VecZnx<&mut [u8]>>, &mut Self){
106        let mut scratch: &mut Self = self;
107        let mut slice: Vec<VecZnx<&mut [u8]>> = Vec::with_capacity(len);
108        for _ in 0..len {
109            let (znx, new_scratch) = scratch.take_vec_znx(n, cols, size);
110            scratch = new_scratch;
111            slice.push(znx);
112        }
113        (slice, scratch)
114    }
115
116    fn take_vmp_pmat<M, B: Backend>(
117        &mut self,
118        module: &M,
119        rows: usize,
120        cols_in: usize,
121        cols_out: usize,
122        size: usize,
123    ) -> (VmpPMat<&mut [u8], B>, &mut Self)
124    where
125        M: VmpPMatBytesOf + ModuleN,
126    {
127        let (take_slice, rem_slice) = self.take_slice(module.bytes_of_vmp_pmat(rows, cols_in, cols_out, size));
128        (
129            VmpPMat::from_data(take_slice, module.n(), rows, cols_in, cols_out, size),
130            rem_slice,
131        )
132    }
133
134    fn take_mat_znx(
135        &mut self,
136        n: usize, 
137        rows: usize,
138        cols_in: usize,
139        cols_out: usize,
140        size: usize,
141    ) -> (MatZnx<&mut [u8]>, &mut Self)
142    {
143        let (take_slice, rem_slice) = self.take_slice(MatZnx::bytes_of(n, rows, cols_in, cols_out, size));
144        (
145            MatZnx::from_data(take_slice, n, rows, cols_in, cols_out, size),
146            rem_slice,
147        )
148    }
149}