1use crate::{
  2    layouts::{Backend, ScalarZnxToRef, Scratch, VecZnxToMut, VecZnxToRef},
  3    source::Source,
  4};
  5
  6pub trait VecZnxNormalizeTmpBytes {
  7    /// Returns the minimum number of bytes necessary for normalization.
  8    fn vec_znx_normalize_tmp_bytes(&self) -> usize;
  9}
 10
 11pub trait VecZnxNormalize<B: Backend> {
 12    /// Normalizes the selected column of `a` and stores the result into the selected column of `res`.
 13    fn vec_znx_normalize<R, A>(&self, basek: usize, res: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch<B>)
 14    where
 15        R: VecZnxToMut,
 16        A: VecZnxToRef;
 17}
 18
 19pub trait VecZnxNormalizeInplace<B: Backend> {
 20    /// Normalizes the selected column of `a`.
 21    fn vec_znx_normalize_inplace<A>(&self, basek: usize, a: &mut A, a_col: usize, scratch: &mut Scratch<B>)
 22    where
 23        A: VecZnxToMut;
 24}
 25
 26pub trait VecZnxAdd {
 27    /// Adds the selected column of `a` to the selected column of `b` and writes the result on the selected column of `res`.
 28    fn vec_znx_add<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
 29    where
 30        R: VecZnxToMut,
 31        A: VecZnxToRef,
 32        B: VecZnxToRef;
 33}
 34
 35pub trait VecZnxAddInplace {
 36    /// Adds the selected column of `a` to the selected column of `res` and writes the result on the selected column of `res`.
 37    fn vec_znx_add_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
 38    where
 39        R: VecZnxToMut,
 40        A: VecZnxToRef;
 41}
 42
 43pub trait VecZnxAddScalar {
 44    /// Adds the selected column of `a` on the selected column and limb of `b` and writes the result on the selected column of `res`.
 45    #[allow(clippy::too_many_arguments)]
 46    fn vec_znx_add_scalar<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize, b_limb: usize)
 47    where
 48        R: VecZnxToMut,
 49        A: ScalarZnxToRef,
 50        B: VecZnxToRef;
 51}
 52
 53pub trait VecZnxAddScalarInplace {
 54    /// Adds the selected column of `a` on the selected column and limb of `res`.
 55    fn vec_znx_add_scalar_inplace<R, A>(&self, res: &mut R, res_col: usize, res_limb: usize, a: &A, a_col: usize)
 56    where
 57        R: VecZnxToMut,
 58        A: ScalarZnxToRef;
 59}
 60
 61pub trait VecZnxSub {
 62    /// Subtracts the selected column of `b` from the selected column of `a` and writes the result on the selected column of `res`.
 63    fn vec_znx_sub<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
 64    where
 65        R: VecZnxToMut,
 66        A: VecZnxToRef,
 67        B: VecZnxToRef;
 68}
 69
 70pub trait VecZnxSubABInplace {
 71    /// Subtracts the selected column of `a` from the selected column of `res` inplace.
 72    ///
 73    /// res\[res_col\] -= a\[a_col\]
 74    fn vec_znx_sub_ab_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
 75    where
 76        R: VecZnxToMut,
 77        A: VecZnxToRef;
 78}
 79
 80pub trait VecZnxSubBAInplace {
 81    /// Subtracts the selected column of `res` from the selected column of `a` and inplace mutates `res`
 82    ///
 83    /// res\[res_col\] = a\[a_col\] - res\[res_col\]
 84    fn vec_znx_sub_ba_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
 85    where
 86        R: VecZnxToMut,
 87        A: VecZnxToRef;
 88}
 89
 90pub trait VecZnxSubScalar {
 91    /// Subtracts the selected column of `a` on the selected column and limb of `b` and writes the result on the selected column of `res`.
 92    #[allow(clippy::too_many_arguments)]
 93    fn vec_znx_sub_scalar<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize, b_limb: usize)
 94    where
 95        R: VecZnxToMut,
 96        A: ScalarZnxToRef,
 97        B: VecZnxToRef;
 98}
 99
100pub trait VecZnxSubScalarInplace {
101    /// Subtracts the selected column of `a` on the selected column and limb of `res`.
102    fn vec_znx_sub_scalar_inplace<R, A>(&self, res: &mut R, res_col: usize, res_limb: usize, a: &A, a_col: usize)
103    where
104        R: VecZnxToMut,
105        A: ScalarZnxToRef;
106}
107
108pub trait VecZnxNegate {
109    // Negates the selected column of `a` and stores the result in `res_col` of `res`.
110    fn vec_znx_negate<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
111    where
112        R: VecZnxToMut,
113        A: VecZnxToRef;
114}
115
116pub trait VecZnxNegateInplace {
117    /// Negates the selected column of `a`.
118    fn vec_znx_negate_inplace<A>(&self, a: &mut A, a_col: usize)
119    where
120        A: VecZnxToMut;
121}
122
123pub trait VecZnxLshTmpBytes {
124    fn vec_znx_lsh_tmp_bytes(&self) -> usize;
125}
126
127pub trait VecZnxLsh<B: Backend> {
128    /// Left shift by k bits all columns of `a`.
129    #[allow(clippy::too_many_arguments)]
130    fn vec_znx_lsh<R, A>(&self, basek: usize, k: usize, r: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch<B>)
131    where
132        R: VecZnxToMut,
133        A: VecZnxToRef;
134}
135
136pub trait VecZnxRshTmpBytes {
137    fn vec_znx_rsh_tmp_bytes(&self) -> usize;
138}
139
140pub trait VecZnxRsh<B: Backend> {
141    /// Right shift by k bits all columns of `a`.
142    #[allow(clippy::too_many_arguments)]
143    fn vec_znx_rsh<R, A>(&self, basek: usize, k: usize, r: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch<B>)
144    where
145        R: VecZnxToMut,
146        A: VecZnxToRef;
147}
148
149pub trait VecZnxLshInplace<B: Backend> {
150    /// Left shift by k bits all columns of `a`.
151    fn vec_znx_lsh_inplace<A>(&self, basek: usize, k: usize, a: &mut A, a_col: usize, scratch: &mut Scratch<B>)
152    where
153        A: VecZnxToMut;
154}
155
156pub trait VecZnxRshInplace<B: Backend> {
157    /// Right shift by k bits all columns of `a`.
158    fn vec_znx_rsh_inplace<A>(&self, basek: usize, k: usize, a: &mut A, a_col: usize, scratch: &mut Scratch<B>)
159    where
160        A: VecZnxToMut;
161}
162
163pub trait VecZnxRotate {
164    /// Multiplies the selected column of `a` by X^k and stores the result in `res_col` of `res`.
165    fn vec_znx_rotate<R, A>(&self, p: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
166    where
167        R: VecZnxToMut,
168        A: VecZnxToRef;
169}
170
171pub trait VecZnxRotateInplaceTmpBytes {
172    fn vec_znx_rotate_inplace_tmp_bytes(&self) -> usize;
173}
174
175pub trait VecZnxRotateInplace<B: Backend> {
176    /// Multiplies the selected column of `a` by X^k.
177    fn vec_znx_rotate_inplace<A>(&self, p: i64, a: &mut A, a_col: usize, scratch: &mut Scratch<B>)
178    where
179        A: VecZnxToMut;
180}
181
182pub trait VecZnxAutomorphism {
183    /// Applies the automorphism X^i -> X^ik on the selected column of `a` and stores the result in `res_col` column of `res`.
184    fn vec_znx_automorphism<R, A>(&self, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
185    where
186        R: VecZnxToMut,
187        A: VecZnxToRef;
188}
189
190pub trait VecZnxAutomorphismInplaceTmpBytes {
191    fn vec_znx_automorphism_inplace_tmp_bytes(&self) -> usize;
192}
193
194pub trait VecZnxAutomorphismInplace<B: Backend> {
195    /// Applies the automorphism X^i -> X^ik on the selected column of `a`.
196    fn vec_znx_automorphism_inplace<R>(&self, k: i64, res: &mut R, res_col: usize, scratch: &mut Scratch<B>)
197    where
198        R: VecZnxToMut;
199}
200
201pub trait VecZnxMulXpMinusOne {
202    fn vec_znx_mul_xp_minus_one<R, A>(&self, p: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
203    where
204        R: VecZnxToMut,
205        A: VecZnxToRef;
206}
207
208pub trait VecZnxMulXpMinusOneInplaceTmpBytes {
209    fn vec_znx_mul_xp_minus_one_inplace_tmp_bytes(&self) -> usize;
210}
211
212pub trait VecZnxMulXpMinusOneInplace<B: Backend> {
213    fn vec_znx_mul_xp_minus_one_inplace<R>(&self, p: i64, res: &mut R, res_col: usize, scratch: &mut Scratch<B>)
214    where
215        R: VecZnxToMut;
216}
217
218pub trait VecZnxSplitRingTmpBytes {
219    fn vec_znx_split_ring_tmp_bytes(&self) -> usize;
220}
221
222pub trait VecZnxSplitRing<B: Backend> {
223    /// Splits the selected columns of `b` into subrings and copies them them into the selected column of `res`.
224    ///
225    /// # Panics
226    ///
227    /// This method requires that all [crate::layouts::VecZnx] of b have the same ring degree
228    /// and that b.n() * b.len() <= a.n()
229    fn vec_znx_split_ring<R, A>(&self, res: &mut [R], res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch<B>)
230    where
231        R: VecZnxToMut,
232        A: VecZnxToRef;
233}
234
235pub trait VecZnxMergeRingsTmpBytes {
236    fn vec_znx_merge_rings_tmp_bytes(&self) -> usize;
237}
238
239pub trait VecZnxMergeRings<B: Backend> {
240    /// Merges the subrings of the selected column of `a` into the selected column of `res`.
241    ///
242    /// # Panics
243    ///
244    /// This method requires that all [crate::layouts::VecZnx] of a have the same ring degree
245    /// and that a.n() * a.len() <= b.n()
246    fn vec_znx_merge_rings<R, A>(&self, res: &mut R, res_col: usize, a: &[A], a_col: usize, scratch: &mut Scratch<B>)
247    where
248        R: VecZnxToMut,
249        A: VecZnxToRef;
250}
251
252pub trait VecZnxSwitchRing {
253    fn vec_znx_switch_ring<R, A>(&self, res: &mut R, res_col: usize, a: &A, col_a: usize)
254    where
255        R: VecZnxToMut,
256        A: VecZnxToRef;
257}
258
259pub trait VecZnxCopy {
260    fn vec_znx_copy<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
261    where
262        R: VecZnxToMut,
263        A: VecZnxToRef;
264}
265
266pub trait VecZnxFillUniform {
267    /// Fills the first `size` size with uniform values in \[-2^{basek-1}, 2^{basek-1}\]
268    fn vec_znx_fill_uniform<R>(&self, basek: usize, res: &mut R, res_col: usize, source: &mut Source)
269    where
270        R: VecZnxToMut;
271}
272
273#[allow(clippy::too_many_arguments)]
274pub trait VecZnxFillNormal {
275    fn vec_znx_fill_normal<R>(
276        &self,
277        basek: usize,
278        res: &mut R,
279        res_col: usize,
280        k: usize,
281        source: &mut Source,
282        sigma: f64,
283        bound: f64,
284    ) where
285        R: VecZnxToMut;
286}
287
288#[allow(clippy::too_many_arguments)]
289pub trait VecZnxAddNormal {
290    /// Adds a discrete normal vector scaled by 2^{-k} with the provided standard deviation and bounded to \[-bound, bound\].
291    fn vec_znx_add_normal<R>(
292        &self,
293        basek: usize,
294        res: &mut R,
295        res_col: usize,
296        k: usize,
297        source: &mut Source,
298        sigma: f64,
299        bound: f64,
300    ) where
301        R: VecZnxToMut;
302}