1use poulpy_hal::{
 2    api::{
 3        ScratchAvailable, VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf, VecZnxIdftApplyConsume, VecZnxNormalize,
 4        VecZnxNormalizeTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
 5    },
 6    layouts::{Backend, DataMut, DataRef, Module, Scratch},
 7};
 8
 9use crate::layouts::{AutomorphismKey, GGLWEInfos, GGSWInfos, GLWESwitchingKey, prepared::GGSWPrepared};
10
11impl AutomorphismKey<Vec<u8>> {
12    pub fn external_product_scratch_space<B: Backend, OUT, IN, GGSW>(
13        module: &Module<B>,
14        out_infos: &OUT,
15        in_infos: &IN,
16        ggsw_infos: &GGSW,
17    ) -> usize
18    where
19        OUT: GGLWEInfos,
20        IN: GGLWEInfos,
21        GGSW: GGSWInfos,
22        Module<B>: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxNormalizeTmpBytes,
23    {
24        GLWESwitchingKey::external_product_scratch_space(module, out_infos, in_infos, ggsw_infos)
25    }
26
27    pub fn external_product_inplace_scratch_space<B: Backend, OUT, GGSW>(
28        module: &Module<B>,
29        out_infos: &OUT,
30        ggsw_infos: &GGSW,
31    ) -> usize
32    where
33        OUT: GGLWEInfos,
34        GGSW: GGSWInfos,
35        Module<B>: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxNormalizeTmpBytes,
36    {
37        GLWESwitchingKey::external_product_inplace_scratch_space(module, out_infos, ggsw_infos)
38    }
39}
40
41impl<DataSelf: DataMut> AutomorphismKey<DataSelf> {
42    pub fn external_product<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
43        &mut self,
44        module: &Module<B>,
45        lhs: &AutomorphismKey<DataLhs>,
46        rhs: &GGSWPrepared<DataRhs, B>,
47        scratch: &mut Scratch<B>,
48    ) where
49        Module<B>: VecZnxDftBytesOf
50            + VmpApplyDftToDftTmpBytes
51            + VecZnxNormalizeTmpBytes
52            + VecZnxDftApply<B>
53            + VmpApplyDftToDft<B>
54            + VmpApplyDftToDftAdd<B>
55            + VecZnxIdftApplyConsume<B>
56            + VecZnxBigNormalize<B>
57            + VecZnxNormalize<B>,
58        Scratch<B>: ScratchAvailable,
59    {
60        self.key.external_product(module, &lhs.key, rhs, scratch);
61    }
62
63    pub fn external_product_inplace<DataRhs: DataRef, B: Backend>(
64        &mut self,
65        module: &Module<B>,
66        rhs: &GGSWPrepared<DataRhs, B>,
67        scratch: &mut Scratch<B>,
68    ) where
69        Module<B>: VecZnxDftBytesOf
70            + VmpApplyDftToDftTmpBytes
71            + VecZnxNormalizeTmpBytes
72            + VecZnxDftApply<B>
73            + VmpApplyDftToDft<B>
74            + VmpApplyDftToDftAdd<B>
75            + VecZnxIdftApplyConsume<B>
76            + VecZnxBigNormalize<B>
77            + VecZnxNormalize<B>,
78        Scratch<B>: ScratchAvailable,
79    {
80        self.key.external_product_inplace(module, rhs, scratch);
81    }
82}