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