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