1use poulpy_hal::{
 2    api::{
 3        DFT, IDFTConsume, ScratchAvailable, TakeVecZnxDft, VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxNormalizeTmpBytes,
 4        VmpApply, VmpApplyAdd, VmpApplyTmpBytes,
 5    },
 6    layouts::{Backend, DataMut, DataRef, Module, Scratch},
 7};
 8
 9use crate::layouts::{GGLWEAutomorphismKey, GGLWESwitchingKey, prepared::GGSWCiphertextPrepared};
10
11impl GGLWEAutomorphismKey<Vec<u8>> {
12    #[allow(clippy::too_many_arguments)]
13    pub fn external_product_scratch_space<B: Backend>(
14        module: &Module<B>,
15        basek: usize,
16        k_out: usize,
17        k_in: usize,
18        ggsw_k: usize,
19        digits: usize,
20        rank: usize,
21    ) -> usize
22    where
23        Module<B>: VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxNormalizeTmpBytes,
24    {
25        GGLWESwitchingKey::external_product_scratch_space(module, basek, k_out, k_in, ggsw_k, digits, rank)
26    }
27
28    pub fn external_product_inplace_scratch_space<B: Backend>(
29        module: &Module<B>,
30        basek: usize,
31        k_out: usize,
32        ggsw_k: usize,
33        digits: usize,
34        rank: usize,
35    ) -> usize
36    where
37        Module<B>: VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxNormalizeTmpBytes,
38    {
39        GGLWESwitchingKey::external_product_inplace_scratch_space(module, basek, k_out, ggsw_k, digits, rank)
40    }
41}
42
43impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
44    pub fn external_product<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
45        &mut self,
46        module: &Module<B>,
47        lhs: &GGLWEAutomorphismKey<DataLhs>,
48        rhs: &GGSWCiphertextPrepared<DataRhs, B>,
49        scratch: &mut Scratch<B>,
50    ) where
51        Module<B>: VecZnxDftAllocBytes
52            + VmpApplyTmpBytes
53            + VecZnxNormalizeTmpBytes
54            + DFT<B>
55            + VmpApply<B>
56            + VmpApplyAdd<B>
57            + IDFTConsume<B>
58            + VecZnxBigNormalize<B>,
59        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
60    {
61        self.key.external_product(module, &lhs.key, rhs, scratch);
62    }
63
64    pub fn external_product_inplace<DataRhs: DataRef, B: Backend>(
65        &mut self,
66        module: &Module<B>,
67        rhs: &GGSWCiphertextPrepared<DataRhs, B>,
68        scratch: &mut Scratch<B>,
69    ) where
70        Module<B>: VecZnxDftAllocBytes
71            + VmpApplyTmpBytes
72            + VecZnxNormalizeTmpBytes
73            + DFT<B>
74            + VmpApply<B>
75            + VmpApplyAdd<B>
76            + IDFTConsume<B>
77            + VecZnxBigNormalize<B>,
78        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
79    {
80        self.key.external_product_inplace(module, rhs, scratch);
81    }
82}