1use backend::hal::{
 2    api::{
 3        ScratchAvailable, TakeVecZnxDft, VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftFromVecZnx,
 4        VecZnxDftToVecZnxBigConsume, VecZnxNormalizeTmpBytes, 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    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        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, n, 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        n: usize,
31        basek: usize,
32        k_out: usize,
33        ggsw_k: usize,
34        digits: usize,
35        rank: usize,
36    ) -> usize
37    where
38        Module<B>: VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxNormalizeTmpBytes,
39    {
40        GGLWESwitchingKey::external_product_inplace_scratch_space(module, n, basek, k_out, ggsw_k, digits, rank)
41    }
42}
43
44impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
45    pub fn external_product<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
46        &mut self,
47        module: &Module<B>,
48        lhs: &GGLWEAutomorphismKey<DataLhs>,
49        rhs: &GGSWCiphertextPrepared<DataRhs, B>,
50        scratch: &mut Scratch<B>,
51    ) where
52        Module<B>: VecZnxDftAllocBytes
53            + VmpApplyTmpBytes
54            + VecZnxNormalizeTmpBytes
55            + VecZnxDftFromVecZnx<B>
56            + VmpApply<B>
57            + VmpApplyAdd<B>
58            + VecZnxDftToVecZnxBigConsume<B>
59            + VecZnxBigNormalize<B>,
60        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
61    {
62        self.key.external_product(module, &lhs.key, rhs, scratch);
63    }
64
65    pub fn external_product_inplace<DataRhs: DataRef, B: Backend>(
66        &mut self,
67        module: &Module<B>,
68        rhs: &GGSWCiphertextPrepared<DataRhs, B>,
69        scratch: &mut Scratch<B>,
70    ) where
71        Module<B>: VecZnxDftAllocBytes
72            + VmpApplyTmpBytes
73            + VecZnxNormalizeTmpBytes
74            + VecZnxDftFromVecZnx<B>
75            + VmpApply<B>
76            + VmpApplyAdd<B>
77            + VecZnxDftToVecZnxBigConsume<B>
78            + VecZnxBigNormalize<B>,
79        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
80    {
81        self.key.external_product_inplace(module, rhs, scratch);
82    }
83}