1use crate::layouts::{
  2    Backend, CnvPVecL, CnvPVecLToMut, CnvPVecLToRef, CnvPVecR, CnvPVecRToMut, CnvPVecRToRef, Scratch, VecZnxBigToMut,
  3    VecZnxDftToMut, VecZnxToRef, ZnxInfos, ZnxViewMut,
  4};
  5
  6pub trait CnvPVecAlloc<BE: Backend> {
  7    fn cnv_pvec_left_alloc(&self, cols: usize, size: usize) -> CnvPVecL<Vec<u8>, BE>;
  8    fn cnv_pvec_right_alloc(&self, cols: usize, size: usize) -> CnvPVecR<Vec<u8>, BE>;
  9}
 10
 11pub trait CnvPVecBytesOf {
 12    fn bytes_of_cnv_pvec_left(&self, cols: usize, size: usize) -> usize;
 13    fn bytes_of_cnv_pvec_right(&self, cols: usize, size: usize) -> usize;
 14}
 15
 16pub trait Convolution<BE: Backend> {
 17    fn cnv_prepare_left_tmp_bytes(&self, res_size: usize, a_size: usize) -> usize;
 18    fn cnv_prepare_left<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
 19    where
 20        R: CnvPVecLToMut<BE> + ZnxInfos + ZnxViewMut<Scalar = BE::ScalarPrep>,
 21        A: VecZnxToRef + ZnxInfos;
 22
 23    fn cnv_prepare_right_tmp_bytes(&self, res_size: usize, a_size: usize) -> usize;
 24    fn cnv_prepare_right<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
 25    where
 26        R: CnvPVecRToMut<BE> + ZnxInfos + ZnxViewMut<Scalar = BE::ScalarPrep>,
 27        A: VecZnxToRef + ZnxInfos;
 28
 29    fn cnv_apply_dft_tmp_bytes(&self, res_size: usize, res_offset: usize, a_size: usize, b_size: usize) -> usize;
 30
 31    fn cnv_by_const_apply_tmp_bytes(&self, res_size: usize, res_offset: usize, a_size: usize, b_size: usize) -> usize;
 32
 33    /// Evaluates a bivariate convolution over Z[X, Y] (x) Z[Y] mod (X^N + 1) where Y = 2^-K over the
 34    /// selected columns and stores the result on the selected column, scaled by 2^{res_offset * Base2K}
 35    ///
 36    /// Behavior is identical to [Convolution::cnv_apply_dft] with `b` treated as a constant polynomial
 37    /// in the X variable, for example:
 38    ///```text
 39    ///       1    X   X^2  X^3
 40    /// a = 1 [a00, a10, a20, a30] = (a00 + a01 * 2^-K) + (a10 + a11 * 2^-K) * X ...
 41    ///     Y [a01, a11, a21, a31]
 42    ///
 43    /// b = 1 [b0] = (b00 + b01 * 2^-K)
 44    ///     Y [b0]
 45    /// ```
 46    /// This method is intended to be used for multiplications by constants that are greater than the base2k.
 47    #[allow(clippy::too_many_arguments)]
 48    fn cnv_by_const_apply<R, A>(
 49        &self,
 50        res: &mut R,
 51        res_offset: usize,
 52        res_col: usize,
 53        a: &A,
 54        a_col: usize,
 55        b: &[i64],
 56        scratch: &mut Scratch<BE>,
 57    ) where
 58        R: VecZnxBigToMut<BE>,
 59        A: VecZnxToRef;
 60
 61    #[allow(clippy::too_many_arguments)]
 62    /// Evaluates a bivariate convolution over Z[X, Y] (x) Z[X, Y] mod (X^N + 1) where Y = 2^-K over the
 63    /// selected columns and stores the result on the selected column, scaled by 2^{res_offset * Base2K}
 64    ///
 65    /// # Example
 66    ///```text
 67    ///       1    X   X^2  X^3
 68    /// a = 1 [a00, a10, a20, a30] = (a00 + a01 * 2^-K) + (a10 + a11 * 2^-K) * X ...
 69    ///     Y [a01, a11, a21, a31]
 70    ///
 71    /// b = 1 [b00, b10, b20, b30] = (b00 + b01 * 2^-K) + (b10 + b11 * 2^-K) * X ...
 72    ///     Y [b01, b11, b21, b31]
 73    ///
 74    /// If res_offset = 0:
 75    ///
 76    ///            1    X   X^2  X^3
 77    /// res = 1  [r00, r10, r20, r30] = (r00 + r01 * 2^-K + r02 * 2^-2K + r03 * 2^-3K) + ... * X + ...
 78    ///       Y  [r01, r11, r21, r31]
 79    ///       Y^2[r02, r12, r22, r32]
 80    ///       Y^3[r03, r13, r23, r33]
 81    ///
 82    /// If res_offset = 1:
 83    ///
 84    ///            1    X   X^2  X^3
 85    /// res = 1  [r01, r11, r21, r31]  = (r01 + r02 * 2^-K + r03 * 2^-2K) + ... * X + ...
 86    ///       Y  [r02, r12, r22, r32]
 87    ///       Y^2[r03, r13, r23, r33]
 88    ///       Y^3[  0,   0,   0 ,  0]
 89    /// ```
 90    /// If res.size() < a.size() + b.size() + k, result is truncated accordingly in the Y dimension.
 91    fn cnv_apply_dft<R, A, B>(
 92        &self,
 93        res: &mut R,
 94        res_offset: usize,
 95        res_col: usize,
 96        a: &A,
 97        a_col: usize,
 98        b: &B,
 99        b_col: usize,
100        scratch: &mut Scratch<BE>,
101    ) where
102        R: VecZnxDftToMut<BE>,
103        A: CnvPVecLToRef<BE>,
104        B: CnvPVecRToRef<BE>;
105
106    fn cnv_pairwise_apply_dft_tmp_bytes(&self, res_size: usize, res_offset: usize, a_size: usize, b_size: usize) -> usize;
107
108    #[allow(clippy::too_many_arguments)]
109    /// Evaluates the bivariate pair-wise convolution res = (a[i] + a[j]) * (b[i] + b[j]).
110    /// If i == j then calls [Convolution::cnv_apply_dft], i.e. res = a[i] * b[i].
111    /// See [Convolution::cnv_apply_dft] for information about the bivariate convolution.
112    fn cnv_pairwise_apply_dft<R, A, B>(
113        &self,
114        res: &mut R,
115        res_offset: usize,
116        res_col: usize,
117        a: &A,
118        b: &B,
119        i: usize,
120        j: usize,
121        scratch: &mut Scratch<BE>,
122    ) where
123        R: VecZnxDftToMut<BE>,
124        A: CnvPVecLToRef<BE>,
125        B: CnvPVecRToRef<BE>;
126}