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