1use crate::hal::{
  2    api::ZnxInfos,
  3    layouts::{Backend, DataRef, MatZnx, ScalarZnx, Scratch, ScratchOwned, SvpPPol, VecZnx, VecZnxBig, VecZnxDft, VmpPMat},
  4};
  5
  6pub unsafe trait ScratchOwnedAllocImpl<B: Backend> {
  7    fn scratch_owned_alloc_impl(size: usize) -> ScratchOwned<B>;
  8}
  9
 10pub unsafe trait ScratchOwnedBorrowImpl<B: Backend> {
 11    fn scratch_owned_borrow_impl(scratch: &mut ScratchOwned<B>) -> &mut Scratch<B>;
 12}
 13
 14pub unsafe trait ScratchFromBytesImpl<B: Backend> {
 15    fn scratch_from_bytes_impl(data: &mut [u8]) -> &mut Scratch<B>;
 16}
 17
 18pub unsafe trait ScratchAvailableImpl<B: Backend> {
 19    fn scratch_available_impl(scratch: &Scratch<B>) -> usize;
 20}
 21
 22pub unsafe trait TakeSliceImpl<B: Backend> {
 23    fn take_slice_impl<T>(scratch: &mut Scratch<B>, len: usize) -> (&mut [T], &mut Scratch<B>);
 24}
 25
 26pub unsafe trait TakeScalarZnxImpl<B: Backend> {
 27    fn take_scalar_znx_impl(scratch: &mut Scratch<B>, n: usize, cols: usize) -> (ScalarZnx<&mut [u8]>, &mut Scratch<B>);
 28}
 29
 30pub unsafe trait TakeSvpPPolImpl<B: Backend> {
 31    fn take_svp_ppol_impl(scratch: &mut Scratch<B>, n: usize, cols: usize) -> (SvpPPol<&mut [u8], B>, &mut Scratch<B>);
 32}
 33
 34pub unsafe trait TakeVecZnxImpl<B: Backend> {
 35    fn take_vec_znx_impl(scratch: &mut Scratch<B>, n: usize, cols: usize, size: usize) -> (VecZnx<&mut [u8]>, &mut Scratch<B>);
 36}
 37
 38pub unsafe trait TakeVecZnxSliceImpl<B: Backend> {
 39    fn take_vec_znx_slice_impl(
 40        scratch: &mut Scratch<B>,
 41        len: usize,
 42        n: usize,
 43        cols: usize,
 44        size: usize,
 45    ) -> (Vec<VecZnx<&mut [u8]>>, &mut Scratch<B>);
 46}
 47
 48pub unsafe trait TakeVecZnxBigImpl<B: Backend> {
 49    fn take_vec_znx_big_impl(
 50        scratch: &mut Scratch<B>,
 51        n: usize,
 52        cols: usize,
 53        size: usize,
 54    ) -> (VecZnxBig<&mut [u8], B>, &mut Scratch<B>);
 55}
 56
 57pub unsafe trait TakeVecZnxDftImpl<B: Backend> {
 58    fn take_vec_znx_dft_impl(
 59        scratch: &mut Scratch<B>,
 60        n: usize,
 61        cols: usize,
 62        size: usize,
 63    ) -> (VecZnxDft<&mut [u8], B>, &mut Scratch<B>);
 64}
 65
 66pub unsafe trait TakeVecZnxDftSliceImpl<B: Backend> {
 67    fn take_vec_znx_dft_slice_impl(
 68        scratch: &mut Scratch<B>,
 69        len: usize,
 70        n: usize,
 71        cols: usize,
 72        size: usize,
 73    ) -> (Vec<VecZnxDft<&mut [u8], B>>, &mut Scratch<B>);
 74}
 75
 76pub unsafe trait TakeVmpPMatImpl<B: Backend> {
 77    fn take_vmp_pmat_impl(
 78        scratch: &mut Scratch<B>,
 79        n: usize,
 80        rows: usize,
 81        cols_in: usize,
 82        cols_out: usize,
 83        size: usize,
 84    ) -> (VmpPMat<&mut [u8], B>, &mut Scratch<B>);
 85}
 86
 87pub unsafe trait TakeMatZnxImpl<B: Backend> {
 88    fn take_mat_znx_impl(
 89        scratch: &mut Scratch<B>,
 90        n: usize,
 91        rows: usize,
 92        cols_in: usize,
 93        cols_out: usize,
 94        size: usize,
 95    ) -> (MatZnx<&mut [u8]>, &mut Scratch<B>);
 96}
 97
 98pub trait TakeLikeImpl<'a, B: Backend, T> {
 99    type Output;
100    fn take_like_impl(scratch: &'a mut Scratch<B>, template: &T) -> (Self::Output, &'a mut Scratch<B>);
101}
102
103impl<'a, B: Backend, D> TakeLikeImpl<'a, B, VmpPMat<D, B>> for B
104where
105    B: TakeVmpPMatImpl<B>,
106    D: DataRef,
107{
108    type Output = VmpPMat<&'a mut [u8], B>;
109
110    fn take_like_impl(scratch: &'a mut Scratch<B>, template: &VmpPMat<D, B>) -> (Self::Output, &'a mut Scratch<B>) {
111        B::take_vmp_pmat_impl(
112            scratch,
113            template.n(),
114            template.rows(),
115            template.cols_in(),
116            template.cols_out(),
117            template.size(),
118        )
119    }
120}
121
122impl<'a, B: Backend, D> TakeLikeImpl<'a, B, MatZnx<D>> for B
123where
124    B: TakeMatZnxImpl<B>,
125    D: DataRef,
126{
127    type Output = MatZnx<&'a mut [u8]>;
128
129    fn take_like_impl(scratch: &'a mut Scratch<B>, template: &MatZnx<D>) -> (Self::Output, &'a mut Scratch<B>) {
130        B::take_mat_znx_impl(
131            scratch,
132            template.n(),
133            template.rows(),
134            template.cols_in(),
135            template.cols_out(),
136            template.size(),
137        )
138    }
139}
140
141impl<'a, B: Backend, D> TakeLikeImpl<'a, B, VecZnxDft<D, B>> for B
142where
143    B: TakeVecZnxDftImpl<B>,
144    D: DataRef,
145{
146    type Output = VecZnxDft<&'a mut [u8], B>;
147
148    fn take_like_impl(scratch: &'a mut Scratch<B>, template: &VecZnxDft<D, B>) -> (Self::Output, &'a mut Scratch<B>) {
149        B::take_vec_znx_dft_impl(scratch, template.n(), template.cols(), template.size())
150    }
151}
152
153impl<'a, B: Backend, D> TakeLikeImpl<'a, B, VecZnxBig<D, B>> for B
154where
155    B: TakeVecZnxBigImpl<B>,
156    D: DataRef,
157{
158    type Output = VecZnxBig<&'a mut [u8], B>;
159
160    fn take_like_impl(scratch: &'a mut Scratch<B>, template: &VecZnxBig<D, B>) -> (Self::Output, &'a mut Scratch<B>) {
161        B::take_vec_znx_big_impl(scratch, template.n(), template.cols(), template.size())
162    }
163}
164
165impl<'a, B: Backend, D> TakeLikeImpl<'a, B, SvpPPol<D, B>> for B
166where
167    B: TakeSvpPPolImpl<B>,
168    D: DataRef,
169{
170    type Output = SvpPPol<&'a mut [u8], B>;
171
172    fn take_like_impl(scratch: &'a mut Scratch<B>, template: &SvpPPol<D, B>) -> (Self::Output, &'a mut Scratch<B>) {
173        B::take_svp_ppol_impl(scratch, template.n(), template.cols())
174    }
175}
176
177impl<'a, B: Backend, D> TakeLikeImpl<'a, B, VecZnx<D>> for B
178where
179    B: TakeVecZnxImpl<B>,
180    D: DataRef,
181{
182    type Output = VecZnx<&'a mut [u8]>;
183
184    fn take_like_impl(scratch: &'a mut Scratch<B>, template: &VecZnx<D>) -> (Self::Output, &'a mut Scratch<B>) {
185        B::take_vec_znx_impl(scratch, template.n(), template.cols(), template.size())
186    }
187}
188
189impl<'a, B: Backend, D> TakeLikeImpl<'a, B, ScalarZnx<D>> for B
190where
191    B: TakeScalarZnxImpl<B>,
192    D: DataRef,
193{
194    type Output = ScalarZnx<&'a mut [u8]>;
195
196    fn take_like_impl(scratch: &'a mut Scratch<B>, template: &ScalarZnx<D>) -> (Self::Output, &'a mut Scratch<B>) {
197        B::take_scalar_znx_impl(scratch, template.n(), template.cols())
198    }
199}