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