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