1use poulpy_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    #[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        ggsw_k: usize,
20        digits: usize,
21        rank: usize,
22    ) -> usize
23    where
24        Module<B>: VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxNormalizeTmpBytes,
25    {
26        GGLWESwitchingKey::external_product_scratch_space(module, n, basek, k_out, k_in, ggsw_k, 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        ggsw_k: usize,
35        digits: usize,
36        rank: usize,
37    ) -> usize
38    where
39        Module<B>: VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxNormalizeTmpBytes,
40    {
41        GGLWESwitchingKey::external_product_inplace_scratch_space(module, n, basek, k_out, ggsw_k, digits, rank)
42    }
43}
44
45impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
46    pub fn external_product<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
47        &mut self,
48        module: &Module<B>,
49        lhs: &GGLWEAutomorphismKey<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        self.key.external_product(module, &lhs.key, rhs, scratch);
64    }
65
66    pub fn external_product_inplace<DataRhs: DataRef, B: Backend>(
67        &mut self,
68        module: &Module<B>,
69        rhs: &GGSWCiphertextPrepared<DataRhs, B>,
70        scratch: &mut Scratch<B>,
71    ) where
72        Module<B>: VecZnxDftAllocBytes
73            + VmpApplyTmpBytes
74            + VecZnxNormalizeTmpBytes
75            + VecZnxDftFromVecZnx<B>
76            + VmpApply<B>
77            + VmpApplyAdd<B>
78            + VecZnxDftToVecZnxBigConsume<B>
79            + VecZnxBigNormalize<B>,
80        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
81    {
82        self.key.external_product_inplace(module, rhs, scratch);
83    }
84}