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