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