1use crate::{
  2    api::{ModuleN, SvpPPolBytesOf, VecZnxBigBytesOf, VecZnxDftBytesOf, VmpPMatBytesOf},
  3    layouts::{Backend, MatZnx, ScalarZnx, Scratch, SvpPPol, VecZnx, VecZnxBig, VecZnxDft, VmpPMat, Zn},
  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 + 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
 41impl<B: Backend> ScratchTakeBasic for Scratch<B> where Self: TakeSlice + ScratchFromBytes<B> {}
 42
 43pub trait ScratchTakeBasic
 44where
 45    Self: TakeSlice,
 46{
 47    fn take_scalar_znx(&mut self, n: usize, cols: usize) -> (ScalarZnx<&mut [u8]>, &mut Self) {
 48        let (take_slice, rem_slice) = self.take_slice(ScalarZnx::bytes_of(n, cols));
 49        (ScalarZnx::from_data(take_slice, n, cols), rem_slice)
 50    }
 51
 52    fn take_svp_ppol<M, B: Backend>(&mut self, module: &M, cols: usize) -> (SvpPPol<&mut [u8], B>, &mut Self)
 53    where
 54        M: SvpPPolBytesOf + ModuleN,
 55    {
 56        let (take_slice, rem_slice) = self.take_slice(module.bytes_of_svp_ppol(cols));
 57        (SvpPPol::from_data(take_slice, module.n(), cols), rem_slice)
 58    }
 59
 60    fn take_zn(&mut self, n: usize, cols: usize, size: usize) -> (Zn<&mut [u8]>, &mut Self) {
 61        let (take_slice, rem_slice) = self.take_slice(Zn::bytes_of(n, cols, size));
 62        (Zn::from_data(take_slice, n, cols, size), rem_slice)
 63    }
 64
 65    fn take_vec_znx(&mut self, n: usize, cols: usize, size: usize) -> (VecZnx<&mut [u8]>, &mut Self) {
 66        let (take_slice, rem_slice) = self.take_slice(VecZnx::bytes_of(n, cols, size));
 67        (VecZnx::from_data(take_slice, n, cols, size), rem_slice)
 68    }
 69
 70    fn take_vec_znx_big<M, B: Backend>(&mut self, module: &M, cols: usize, size: usize) -> (VecZnxBig<&mut [u8], B>, &mut Self)
 71    where
 72        M: VecZnxBigBytesOf + ModuleN,
 73    {
 74        let (take_slice, rem_slice) = self.take_slice(module.bytes_of_vec_znx_big(cols, size));
 75        (
 76            VecZnxBig::from_data(take_slice, module.n(), cols, size),
 77            rem_slice,
 78        )
 79    }
 80
 81    fn take_vec_znx_dft<M, B: Backend>(&mut self, module: &M, cols: usize, size: usize) -> (VecZnxDft<&mut [u8], B>, &mut Self)
 82    where
 83        M: VecZnxDftBytesOf + ModuleN,
 84    {
 85        let (take_slice, rem_slice) = self.take_slice(module.bytes_of_vec_znx_dft(cols, size));
 86
 87        (
 88            VecZnxDft::from_data(take_slice, module.n(), cols, size),
 89            rem_slice,
 90        )
 91    }
 92
 93    fn take_vec_znx_dft_slice<M, B: Backend>(
 94        &mut self,
 95        module: &M,
 96        len: usize,
 97        cols: usize,
 98        size: usize,
 99    ) -> (Vec<VecZnxDft<&mut [u8], B>>, &mut Self)
100    where
101        M: VecZnxDftBytesOf + ModuleN,
102    {
103        let mut scratch: &mut Self = self;
104        let mut slice: Vec<VecZnxDft<&mut [u8], B>> = Vec::with_capacity(len);
105        for _ in 0..len {
106            let (znx, new_scratch) = scratch.take_vec_znx_dft(module, cols, size);
107            scratch = new_scratch;
108            slice.push(znx);
109        }
110        (slice, scratch)
111    }
112
113    fn take_vec_znx_slice(&mut self, len: usize, n: usize, cols: usize, size: usize) -> (Vec<VecZnx<&mut [u8]>>, &mut Self) {
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(n, 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(
143        &mut self,
144        n: usize,
145        rows: usize,
146        cols_in: usize,
147        cols_out: usize,
148        size: usize,
149    ) -> (MatZnx<&mut [u8]>, &mut Self) {
150        let (take_slice, rem_slice) = self.take_slice(MatZnx::bytes_of(n, rows, cols_in, cols_out, size));
151        (
152            MatZnx::from_data(take_slice, n, rows, cols_in, cols_out, size),
153            rem_slice,
154        )
155    }
156}