1use crate::{
  2    api::{
  3        ModuleN, ScratchTakeBasic, SvpApplyDftToDft, SvpPPolAlloc, SvpPPolBytesOf, SvpPrepare, VecZnxDftAddScaledInplace,
  4        VecZnxDftBytesOf, VecZnxDftZero,
  5    },
  6    layouts::{Backend, Module, Scratch, VecZnxDftToMut, VecZnxDftToRef, VecZnxToRef, ZnxInfos},
  7};
  8
  9impl<BE: Backend> Convolution<BE> for Module<BE>
 10where
 11    Self: Sized
 12        + ModuleN
 13        + SvpPPolAlloc<BE>
 14        + SvpApplyDftToDft<BE>
 15        + SvpPrepare<BE>
 16        + SvpPPolBytesOf
 17        + VecZnxDftBytesOf
 18        + VecZnxDftAddScaledInplace<BE>
 19        + VecZnxDftZero<BE>,
 20    Scratch<BE>: ScratchTakeBasic,
 21{
 22}
 23
 24pub trait Convolution<BE: Backend>
 25where
 26    Self: Sized
 27        + ModuleN
 28        + SvpPPolAlloc<BE>
 29        + SvpApplyDftToDft<BE>
 30        + SvpPrepare<BE>
 31        + SvpPPolBytesOf
 32        + VecZnxDftBytesOf
 33        + VecZnxDftAddScaledInplace<BE>
 34        + VecZnxDftZero<BE>,
 35    Scratch<BE>: ScratchTakeBasic,
 36{
 37    fn convolution_tmp_bytes(&self, b_size: usize) -> usize {
 38        self.bytes_of_svp_ppol(1) + self.bytes_of_vec_znx_dft(1, b_size)
 39    }
 40
 41    fn bivariate_convolution_full<R, A, B>(&self, k: i64, res: &mut R, a: &A, b: &B, scratch: &mut Scratch<BE>)
 42    where
 43        R: VecZnxDftToMut<BE>,
 44        A: VecZnxToRef,
 45        B: VecZnxDftToRef<BE>,
 46    {
 47        let res: &mut crate::layouts::VecZnxDft<&mut [u8], BE> = &mut res.to_mut();
 48        let a: &crate::layouts::VecZnx<&[u8]> = &a.to_ref();
 49        let b: &crate::layouts::VecZnxDft<&[u8], BE> = &b.to_ref();
 50
 51        let res_cols: usize = res.cols();
 52        let a_cols: usize = a.cols();
 53        let b_cols: usize = b.cols();
 54
 55        assert!(res_cols >= a_cols + b_cols - 1);
 56
 57        for res_col in 0..res_cols {
 58            let a_min: usize = res_col.saturating_sub(b_cols - 1);
 59            let a_max: usize = res_col.min(a_cols - 1);
 60            self.bivariate_convolution_single(k, res, res_col, a, a_min, b, res_col - a_min, scratch);
 61            for a_col in a_min + 1..a_max + 1 {
 62                self.bivariate_convolution_single_add(k, res, res_col, a, a_col, b, res_col - a_col, scratch);
 63            }
 64        }
 65    }
 66
 67    /// Evaluates a bivariate convolution over Z[X, Y] / (X^N + 1) where Y = 2^-K over the
 68    /// selected columsn and stores the result on the selected column, scaled by 2^{k * Base2K}
 69    ///
 70    /// # Example
 71    /// a = [a00, a10, a20, a30] = (a00 * 2^-K + a01 * 2^-2K) + (a10 * 2^-K + a11 * 2^-2K) * X ...
 72    ///     [a01, a11, a21, a31]
 73    ///
 74    /// b = [b00, b10, b20, b30] = (b00 * 2^-K + b01 * 2^-2K) + (b10 * 2^-K + b11 * 2^-2K) * X ...
 75    ///     [b01, b11, b21, b31]
 76    ///
 77    /// If k = 0:
 78    /// res = [  0,   0,   0,   0] = (r01 * 2^-2K + r02 * 2^-3K + r03 * 2^-4K + r04 * 2^-5K) + ...
 79    ///       [r01, r11, r21, r31]
 80    ///       [r02, r12, r22, r32]
 81    ///       [r03, r13, r23, r33]
 82    ///       [r04, r14, r24, r34]
 83    ///
 84    /// If k = 1:
 85    /// res = [r01, r11, r21, r31] = (r01 * 2^-K + r02 * 2^-2K + r03 * 2^-3K + r04 * 2^-4K + r05 * 2^-5K) + ...
 86    ///       [r02, r12, r22, r32]
 87    ///       [r03, r13, r23, r33]
 88    ///       [r04, r14, r24, r34]
 89    ///       [r05, r15, r25, r35]
 90    ///
 91    /// If k = -1:
 92    /// res = [  0,   0,   0,   0] = (r01 * 2^-3K + r02 * 2^-4K + r03 * 2^-5K) + ...
 93    ///       [  0,   0,   0,   0]
 94    ///       [r01, r11, r21, r31]
 95    ///       [r02, r12, r22, r32]
 96    ///       [r03, r13, r23, r33]
 97    ///
 98    /// If res.size() < a.size() + b.size() + 1 + k, result is truncated accordingly in the Y dimension.
 99    fn bivariate_convolution_single_add<R, A, B>(
100        &self,
101        k: i64,
102        res: &mut R,
103        res_col: usize,
104        a: &A,
105        a_col: usize,
106        b: &B,
107        b_col: usize,
108        scratch: &mut Scratch<BE>,
109    ) where
110        R: VecZnxDftToMut<BE>,
111        A: VecZnxToRef,
112        B: VecZnxDftToRef<BE>,
113    {
114        let res: &mut crate::layouts::VecZnxDft<&mut [u8], BE> = &mut res.to_mut();
115        let a: &crate::layouts::VecZnx<&[u8]> = &a.to_ref();
116        let b: &crate::layouts::VecZnxDft<&[u8], BE> = &b.to_ref();
117
118        let (mut ppol, scratch_1) = scratch.take_svp_ppol(self, 1);
119        let (mut res_tmp, _) = scratch_1.take_vec_znx_dft(self, 1, b.size());
120
121        for a_limb in 0..a.size() {
122            self.svp_prepare(&mut ppol, 0, &a.as_scalar_znx_ref(a_col, a_limb), 0);
123            self.svp_apply_dft_to_dft(&mut res_tmp, 0, &ppol, 0, b, b_col);
124            self.vec_znx_dft_add_scaled_inplace(res, res_col, &res_tmp, 0, -(1 + a_limb as i64) + k);
125        }
126    }
127
128    fn bivariate_convolution_single<R, A, B>(
129        &self,
130        k: i64,
131        res: &mut R,
132        res_col: usize,
133        a: &A,
134        a_col: usize,
135        b: &B,
136        b_col: usize,
137        scratch: &mut Scratch<BE>,
138    ) where
139        R: VecZnxDftToMut<BE>,
140        A: VecZnxToRef,
141        B: VecZnxDftToRef<BE>,
142    {
143        self.vec_znx_dft_zero(res, res_col);
144        self.bivariate_convolution_single_add(k, res, res_col, a, a_col, b, b_col, scratch);
145    }
146}