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