1use crate::hal::{
  2    api::{
  3        ScratchAvailable, ScratchFromBytes, ScratchOwnedAlloc, ScratchOwnedBorrow, TakeLike, TakeMatZnx, TakeScalarZnx,
  4        TakeSlice, TakeSvpPPol, TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, TakeVecZnxDftSlice, TakeVecZnxSlice, TakeVmpPMat,
  5    },
  6    layouts::{
  7        Backend, DataRef, MatZnx, Module, ScalarZnx, Scratch, ScratchOwned, SvpPPol, VecZnx, VecZnxBig, VecZnxDft, VmpPMat,
  8    },
  9    oep::{
 10        ScratchAvailableImpl, ScratchFromBytesImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeLikeImpl, TakeMatZnxImpl,
 11        TakeScalarZnxImpl, TakeSliceImpl, TakeSvpPPolImpl, TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxDftSliceImpl,
 12        TakeVecZnxImpl, TakeVecZnxSliceImpl, TakeVmpPMatImpl,
 13    },
 14};
 15
 16impl<B> ScratchOwnedAlloc<B> for ScratchOwned<B>
 17where
 18    B: Backend + ScratchOwnedAllocImpl<B>,
 19{
 20    fn alloc(size: usize) -> Self {
 21        B::scratch_owned_alloc_impl(size)
 22    }
 23}
 24
 25impl<B> ScratchOwnedBorrow<B> for ScratchOwned<B>
 26where
 27    B: Backend + ScratchOwnedBorrowImpl<B>,
 28{
 29    fn borrow(&mut self) -> &mut Scratch<B> {
 30        B::scratch_owned_borrow_impl(self)
 31    }
 32}
 33
 34impl<B> ScratchFromBytes<B> for Scratch<B>
 35where
 36    B: Backend + ScratchFromBytesImpl<B>,
 37{
 38    fn from_bytes(data: &mut [u8]) -> &mut Scratch<B> {
 39        B::scratch_from_bytes_impl(data)
 40    }
 41}
 42
 43impl<B> ScratchAvailable for Scratch<B>
 44where
 45    B: Backend + ScratchAvailableImpl<B>,
 46{
 47    fn available(&self) -> usize {
 48        B::scratch_available_impl(self)
 49    }
 50}
 51
 52impl<B> TakeSlice for Scratch<B>
 53where
 54    B: Backend + TakeSliceImpl<B>,
 55{
 56    fn take_slice<T>(&mut self, len: usize) -> (&mut [T], &mut Self) {
 57        B::take_slice_impl(self, len)
 58    }
 59}
 60
 61impl<B> TakeScalarZnx<B> for Scratch<B>
 62where
 63    B: Backend + TakeScalarZnxImpl<B>,
 64{
 65    fn take_scalar_znx(&mut self, module: &Module<B>, cols: usize) -> (ScalarZnx<&mut [u8]>, &mut Self) {
 66        B::take_scalar_znx_impl(self, module.n(), cols)
 67    }
 68}
 69
 70impl<B> TakeSvpPPol<B> for Scratch<B>
 71where
 72    B: Backend + TakeSvpPPolImpl<B>,
 73{
 74    fn take_svp_ppol(&mut self, module: &Module<B>, cols: usize) -> (SvpPPol<&mut [u8], B>, &mut Self) {
 75        B::take_svp_ppol_impl(self, module.n(), cols)
 76    }
 77}
 78
 79impl<B> TakeVecZnx<B> for Scratch<B>
 80where
 81    B: Backend + TakeVecZnxImpl<B>,
 82{
 83    fn take_vec_znx(&mut self, module: &Module<B>, cols: usize, size: usize) -> (VecZnx<&mut [u8]>, &mut Self) {
 84        B::take_vec_znx_impl(self, module.n(), cols, size)
 85    }
 86}
 87
 88impl<B> TakeVecZnxSlice<B> for Scratch<B>
 89where
 90    B: Backend + TakeVecZnxSliceImpl<B>,
 91{
 92    fn take_vec_znx_slice(
 93        &mut self,
 94        len: usize,
 95        module: &Module<B>,
 96        cols: usize,
 97        size: usize,
 98    ) -> (Vec<VecZnx<&mut [u8]>>, &mut Self) {
 99        B::take_vec_znx_slice_impl(self, len, module.n(), cols, size)
100    }
101}
102
103impl<B> TakeVecZnxBig<B> for Scratch<B>
104where
105    B: Backend + TakeVecZnxBigImpl<B>,
106{
107    fn take_vec_znx_big(&mut self, module: &Module<B>, cols: usize, size: usize) -> (VecZnxBig<&mut [u8], B>, &mut Self) {
108        B::take_vec_znx_big_impl(self, module.n(), cols, size)
109    }
110}
111
112impl<B> TakeVecZnxDft<B> for Scratch<B>
113where
114    B: Backend + TakeVecZnxDftImpl<B>,
115{
116    fn take_vec_znx_dft(&mut self, module: &Module<B>, cols: usize, size: usize) -> (VecZnxDft<&mut [u8], B>, &mut Self) {
117        B::take_vec_znx_dft_impl(self, module.n(), cols, size)
118    }
119}
120
121impl<B> TakeVecZnxDftSlice<B> for Scratch<B>
122where
123    B: Backend + TakeVecZnxDftSliceImpl<B>,
124{
125    fn take_vec_znx_dft_slice(
126        &mut self,
127        len: usize,
128        module: &Module<B>,
129        cols: usize,
130        size: usize,
131    ) -> (Vec<VecZnxDft<&mut [u8], B>>, &mut Self) {
132        B::take_vec_znx_dft_slice_impl(self, len, module.n(), cols, size)
133    }
134}
135
136impl<B> TakeVmpPMat<B> for Scratch<B>
137where
138    B: Backend + TakeVmpPMatImpl<B>,
139{
140    fn take_vmp_pmat(
141        &mut self,
142        module: &Module<B>,
143        rows: usize,
144        cols_in: usize,
145        cols_out: usize,
146        size: usize,
147    ) -> (VmpPMat<&mut [u8], B>, &mut Self) {
148        B::take_vmp_pmat_impl(self, module.n(), rows, cols_in, cols_out, size)
149    }
150}
151
152impl<B> TakeMatZnx<B> for Scratch<B>
153where
154    B: Backend + TakeMatZnxImpl<B>,
155{
156    fn take_mat_znx(
157        &mut self,
158        module: &Module<B>,
159        rows: usize,
160        cols_in: usize,
161        cols_out: usize,
162        size: usize,
163    ) -> (MatZnx<&mut [u8]>, &mut Self) {
164        B::take_mat_znx_impl(self, module.n(), rows, cols_in, cols_out, size)
165    }
166}
167
168impl<'a, B: Backend, D> TakeLike<'a, B, ScalarZnx<D>> for Scratch<B>
169where
170    B: TakeLikeImpl<'a, B, ScalarZnx<D>, Output = ScalarZnx<&'a mut [u8]>>,
171    D: DataRef,
172{
173    type Output = ScalarZnx<&'a mut [u8]>;
174    fn take_like(&'a mut self, template: &ScalarZnx<D>) -> (Self::Output, &'a mut Self) {
175        B::take_like_impl(self, template)
176    }
177}
178
179impl<'a, B: Backend, D> TakeLike<'a, B, SvpPPol<D, B>> for Scratch<B>
180where
181    B: TakeLikeImpl<'a, B, SvpPPol<D, B>, Output = SvpPPol<&'a mut [u8], B>>,
182    D: DataRef,
183{
184    type Output = SvpPPol<&'a mut [u8], B>;
185    fn take_like(&'a mut self, template: &SvpPPol<D, B>) -> (Self::Output, &'a mut Self) {
186        B::take_like_impl(self, template)
187    }
188}
189
190impl<'a, B: Backend, D> TakeLike<'a, B, VecZnx<D>> for Scratch<B>
191where
192    B: TakeLikeImpl<'a, B, VecZnx<D>, Output = VecZnx<&'a mut [u8]>>,
193    D: DataRef,
194{
195    type Output = VecZnx<&'a mut [u8]>;
196    fn take_like(&'a mut self, template: &VecZnx<D>) -> (Self::Output, &'a mut Self) {
197        B::take_like_impl(self, template)
198    }
199}
200
201impl<'a, B: Backend, D> TakeLike<'a, B, VecZnxBig<D, B>> for Scratch<B>
202where
203    B: TakeLikeImpl<'a, B, VecZnxBig<D, B>, Output = VecZnxBig<&'a mut [u8], B>>,
204    D: DataRef,
205{
206    type Output = VecZnxBig<&'a mut [u8], B>;
207    fn take_like(&'a mut self, template: &VecZnxBig<D, B>) -> (Self::Output, &'a mut Self) {
208        B::take_like_impl(self, template)
209    }
210}
211
212impl<'a, B: Backend, D> TakeLike<'a, B, VecZnxDft<D, B>> for Scratch<B>
213where
214    B: TakeLikeImpl<'a, B, VecZnxDft<D, B>, Output = VecZnxDft<&'a mut [u8], B>>,
215    D: DataRef,
216{
217    type Output = VecZnxDft<&'a mut [u8], B>;
218    fn take_like(&'a mut self, template: &VecZnxDft<D, B>) -> (Self::Output, &'a mut Self) {
219        B::take_like_impl(self, template)
220    }
221}
222
223impl<'a, B: Backend, D> TakeLike<'a, B, MatZnx<D>> for Scratch<B>
224where
225    B: TakeLikeImpl<'a, B, MatZnx<D>, Output = MatZnx<&'a mut [u8]>>,
226    D: DataRef,
227{
228    type Output = MatZnx<&'a mut [u8]>;
229    fn take_like(&'a mut self, template: &MatZnx<D>) -> (Self::Output, &'a mut Self) {
230        B::take_like_impl(self, template)
231    }
232}
233
234impl<'a, B: Backend, D> TakeLike<'a, B, VmpPMat<D, B>> for Scratch<B>
235where
236    B: TakeLikeImpl<'a, B, VmpPMat<D, B>, Output = VmpPMat<&'a mut [u8], B>>,
237    D: DataRef,
238{
239    type Output = VmpPMat<&'a mut [u8], B>;
240    fn take_like(&'a mut self, template: &VmpPMat<D, B>) -> (Self::Output, &'a mut Self) {
241        B::take_like_impl(self, template)
242    }
243}