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