1use backend::hal::{
  2    api::{ScratchAvailable, TakeVecZnxDft, ZnxZero},
  3    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  4};
  5
  6use crate::{GGSWCiphertext, GGSWCiphertextExec, GLWECiphertext, GLWEExternalProductFamily, Infos};
  7
  8impl GGSWCiphertext<Vec<u8>> {
  9    pub fn external_product_scratch_space<B: Backend>(
 10        module: &Module<B>,
 11        n: usize,
 12        basek: usize,
 13        k_out: usize,
 14        k_in: usize,
 15        k_ggsw: usize,
 16        digits: usize,
 17        rank: usize,
 18    ) -> usize
 19    where
 20        Module<B>: GLWEExternalProductFamily<B>,
 21    {
 22        GLWECiphertext::external_product_scratch_space(module, n, basek, k_out, k_in, k_ggsw, digits, rank)
 23    }
 24
 25    pub fn external_product_inplace_scratch_space<B: Backend>(
 26        module: &Module<B>,
 27        n: usize,
 28        basek: usize,
 29        k_out: usize,
 30        k_ggsw: usize,
 31        digits: usize,
 32        rank: usize,
 33    ) -> usize
 34    where
 35        Module<B>: GLWEExternalProductFamily<B>,
 36    {
 37        GLWECiphertext::external_product_inplace_scratch_space(module, n, basek, k_out, k_ggsw, digits, rank)
 38    }
 39}
 40
 41impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
 42    pub fn external_product<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
 43        &mut self,
 44        module: &Module<B>,
 45        lhs: &GGSWCiphertext<DataLhs>,
 46        rhs: &GGSWCiphertextExec<DataRhs, B>,
 47        scratch: &mut Scratch<B>,
 48    ) where
 49        Module<B>: GLWEExternalProductFamily<B>,
 50        Scratch<B>: ScratchAvailable + TakeVecZnxDft<B>,
 51    {
 52        #[cfg(debug_assertions)]
 53        {
 54            use crate::{GGSWCiphertext, Infos};
 55
 56            assert_eq!(lhs.n(), self.n());
 57            assert_eq!(rhs.n(), self.n());
 58
 59            assert_eq!(
 60                self.rank(),
 61                lhs.rank(),
 62                "ggsw_out rank: {} != ggsw_in rank: {}",
 63                self.rank(),
 64                lhs.rank()
 65            );
 66            assert_eq!(
 67                self.rank(),
 68                rhs.rank(),
 69                "ggsw_in rank: {} != ggsw_apply rank: {}",
 70                self.rank(),
 71                rhs.rank()
 72            );
 73
 74            assert!(
 75                scratch.available()
 76                    >= GGSWCiphertext::external_product_scratch_space(
 77                        module,
 78                        self.n(),
 79                        self.basek(),
 80                        self.k(),
 81                        lhs.k(),
 82                        rhs.k(),
 83                        rhs.digits(),
 84                        rhs.rank()
 85                    )
 86            )
 87        }
 88
 89        let min_rows: usize = self.rows().min(lhs.rows());
 90
 91        (0..self.rank() + 1).for_each(|col_i| {
 92            (0..min_rows).for_each(|row_j| {
 93                self.at_mut(row_j, col_i)
 94                    .external_product(module, &lhs.at(row_j, col_i), rhs, scratch);
 95            });
 96            (min_rows..self.rows()).for_each(|row_i| {
 97                self.at_mut(row_i, col_i).data.zero();
 98            });
 99        });
100    }
101
102    pub fn external_product_inplace<DataRhs: DataRef, B: Backend>(
103        &mut self,
104        module: &Module<B>,
105        rhs: &GGSWCiphertextExec<DataRhs, B>,
106        scratch: &mut Scratch<B>,
107    ) where
108        Module<B>: GLWEExternalProductFamily<B>,
109        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
110    {
111        #[cfg(debug_assertions)]
112        {
113            assert_eq!(rhs.n(), self.n());
114            assert_eq!(
115                self.rank(),
116                rhs.rank(),
117                "ggsw_out rank: {} != ggsw_apply: {}",
118                self.rank(),
119                rhs.rank()
120            );
121        }
122
123        (0..self.rank() + 1).for_each(|col_i| {
124            (0..self.rows()).for_each(|row_j| {
125                self.at_mut(row_j, col_i)
126                    .external_product_inplace(module, rhs, scratch);
127            });
128        });
129    }
130}