1use rand_distr::Distribution;
  2
  3use crate::{
  4    layouts::{Backend, Scratch, VecZnxBigOwned, VecZnxBigToMut, VecZnxBigToRef, VecZnxToMut, VecZnxToRef},
  5    source::Source,
  6};
  7
  8/// Allocates as [crate::layouts::VecZnxBig].
  9pub trait VecZnxBigAlloc<B: Backend> {
 10    fn vec_znx_big_alloc(&self, cols: usize, size: usize) -> VecZnxBigOwned<B>;
 11}
 12
 13/// Returns the size in bytes to allocate a [crate::layouts::VecZnxBig].
 14pub trait VecZnxBigAllocBytes {
 15    fn vec_znx_big_alloc_bytes(&self, cols: usize, size: usize) -> usize;
 16}
 17
 18/// Consume a vector of bytes into a [crate::layouts::VecZnxBig].
 19/// User must ensure that bytes is memory aligned and that it length is equal to [VecZnxBigAllocBytes].
 20pub trait VecZnxBigFromBytes<B: Backend> {
 21    fn vec_znx_big_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxBigOwned<B>;
 22}
 23
 24#[allow(clippy::too_many_arguments)]
 25/// Add a discrete normal distribution on res.
 26///
 27/// # Arguments
 28/// * `basek`: base two logarithm of the bivariate representation
 29/// * `res`: receiver.
 30/// * `res_col`: column of the receiver on which the operation is performed/stored.
 31/// * `k`:
 32/// * `source`: random coin source.
 33/// * `sigma`: standard deviation of the discrete normal distribution.
 34/// * `bound`: rejection sampling bound.
 35pub trait VecZnxBigAddNormal<B: Backend> {
 36    fn vec_znx_big_add_normal<R: VecZnxBigToMut<B>>(
 37        &self,
 38        basek: usize,
 39        res: &mut R,
 40        res_col: usize,
 41        k: usize,
 42        source: &mut Source,
 43        sigma: f64,
 44        bound: f64,
 45    );
 46}
 47
 48#[allow(clippy::too_many_arguments)]
 49pub trait VecZnxBigFillNormal<B: Backend> {
 50    fn vec_znx_big_fill_normal<R: VecZnxBigToMut<B>>(
 51        &self,
 52        basek: usize,
 53        res: &mut R,
 54        res_col: usize,
 55        k: usize,
 56        source: &mut Source,
 57        sigma: f64,
 58        bound: f64,
 59    );
 60}
 61
 62#[allow(clippy::too_many_arguments)]
 63pub trait VecZnxBigFillDistF64<B: Backend> {
 64    fn vec_znx_big_fill_dist_f64<R: VecZnxBigToMut<B>, D: Distribution<f64>>(
 65        &self,
 66        basek: usize,
 67        res: &mut R,
 68        res_col: usize,
 69        k: usize,
 70        source: &mut Source,
 71        dist: D,
 72        bound: f64,
 73    );
 74}
 75
 76#[allow(clippy::too_many_arguments)]
 77pub trait VecZnxBigAddDistF64<B: Backend> {
 78    fn vec_znx_big_add_dist_f64<R: VecZnxBigToMut<B>, D: Distribution<f64>>(
 79        &self,
 80        basek: usize,
 81        res: &mut R,
 82        res_col: usize,
 83        k: usize,
 84        source: &mut Source,
 85        dist: D,
 86        bound: f64,
 87    );
 88}
 89
 90pub trait VecZnxBigAdd<B: Backend> {
 91    /// Adds `a` to `b` and stores the result on `c`.
 92    fn vec_znx_big_add<R, A, C>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &C, b_col: usize)
 93    where
 94        R: VecZnxBigToMut<B>,
 95        A: VecZnxBigToRef<B>,
 96        C: VecZnxBigToRef<B>;
 97}
 98
 99pub trait VecZnxBigAddInplace<B: Backend> {
100    /// Adds `a` to `b` and stores the result on `b`.
101    fn vec_znx_big_add_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
102    where
103        R: VecZnxBigToMut<B>,
104        A: VecZnxBigToRef<B>;
105}
106
107pub trait VecZnxBigAddSmall<B: Backend> {
108    /// Adds `a` to `b` and stores the result on `c`.
109    fn vec_znx_big_add_small<R, A, C>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &C, b_col: usize)
110    where
111        R: VecZnxBigToMut<B>,
112        A: VecZnxBigToRef<B>,
113        C: VecZnxToRef;
114}
115
116pub trait VecZnxBigAddSmallInplace<B: Backend> {
117    /// Adds `a` to `b` and stores the result on `b`.
118    fn vec_znx_big_add_small_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
119    where
120        R: VecZnxBigToMut<B>,
121        A: VecZnxToRef;
122}
123
124pub trait VecZnxBigSub<B: Backend> {
125    /// Subtracts `a` to `b` and stores the result on `c`.
126    fn vec_znx_big_sub<R, A, C>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &C, b_col: usize)
127    where
128        R: VecZnxBigToMut<B>,
129        A: VecZnxBigToRef<B>,
130        C: VecZnxBigToRef<B>;
131}
132
133pub trait VecZnxBigSubABInplace<B: Backend> {
134    /// Subtracts `a` from `b` and stores the result on `b`.
135    fn vec_znx_big_sub_ab_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
136    where
137        R: VecZnxBigToMut<B>,
138        A: VecZnxBigToRef<B>;
139}
140
141pub trait VecZnxBigSubBAInplace<B: Backend> {
142    /// Subtracts `b` from `a` and stores the result on `b`.
143    fn vec_znx_big_sub_ba_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
144    where
145        R: VecZnxBigToMut<B>,
146        A: VecZnxBigToRef<B>;
147}
148
149pub trait VecZnxBigSubSmallA<B: Backend> {
150    /// Subtracts `b` from `a` and stores the result on `c`.
151    fn vec_znx_big_sub_small_a<R, A, C>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &C, b_col: usize)
152    where
153        R: VecZnxBigToMut<B>,
154        A: VecZnxToRef,
155        C: VecZnxBigToRef<B>;
156}
157
158pub trait VecZnxBigSubSmallAInplace<B: Backend> {
159    /// Subtracts `a` from `res` and stores the result on `res`.
160    fn vec_znx_big_sub_small_a_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
161    where
162        R: VecZnxBigToMut<B>,
163        A: VecZnxToRef;
164}
165
166pub trait VecZnxBigSubSmallB<B: Backend> {
167    /// Subtracts `b` from `a` and stores the result on `c`.
168    fn vec_znx_big_sub_small_b<R, A, C>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &C, b_col: usize)
169    where
170        R: VecZnxBigToMut<B>,
171        A: VecZnxBigToRef<B>,
172        C: VecZnxToRef;
173}
174
175pub trait VecZnxBigSubSmallBInplace<B: Backend> {
176    /// Subtracts `res` from `a` and stores the result on `res`.
177    fn vec_znx_big_sub_small_b_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
178    where
179        R: VecZnxBigToMut<B>,
180        A: VecZnxToRef;
181}
182
183pub trait VecZnxBigNegateInplace<B: Backend> {
184    fn vec_znx_big_negate_inplace<A>(&self, a: &mut A, a_col: usize)
185    where
186        A: VecZnxBigToMut<B>;
187}
188
189pub trait VecZnxBigNormalizeTmpBytes {
190    fn vec_znx_big_normalize_tmp_bytes(&self) -> usize;
191}
192
193pub trait VecZnxBigNormalize<B: Backend> {
194    fn vec_znx_big_normalize<R, A>(
195        &self,
196        basek: usize,
197        res: &mut R,
198        res_col: usize,
199        a: &A,
200        a_col: usize,
201        scratch: &mut Scratch<B>,
202    ) where
203        R: VecZnxToMut,
204        A: VecZnxBigToRef<B>;
205}
206
207pub trait VecZnxBigAutomorphism<B: Backend> {
208    /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `b`.
209    fn vec_znx_big_automorphism<R, A>(&self, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
210    where
211        R: VecZnxBigToMut<B>,
212        A: VecZnxBigToRef<B>;
213}
214
215pub trait VecZnxBigAutomorphismInplace<B: Backend> {
216    /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `a`.
217    fn vec_znx_big_automorphism_inplace<A>(&self, k: i64, a: &mut A, a_col: usize)
218    where
219        A: VecZnxBigToMut<B>;
220}