1use poulpy_hal::{
  2    api::{
  3        ScratchAvailable, TakeVecZnxBig, TakeVecZnxDft, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace, VecZnxBigAllocBytes,
  4        VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftAddInplace, VecZnxDftAllocBytes, VecZnxDftCopy,
  5        VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VecZnxDftToVecZnxBigTmpA, VecZnxNormalizeTmpBytes, VmpApply,
  6        VmpApplyAdd, VmpApplyTmpBytes,
  7    },
  8    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  9};
 10
 11use crate::layouts::{
 12    GGSWCiphertext, GLWECiphertext, Infos,
 13    prepared::{GGLWEAutomorphismKeyPrepared, GGLWETensorKeyPrepared},
 14};
 15
 16impl GGSWCiphertext<Vec<u8>> {
 17    #[allow(clippy::too_many_arguments)]
 18    pub fn automorphism_scratch_space<B: Backend>(
 19        module: &Module<B>,
 20        basek: usize,
 21        k_out: usize,
 22        k_in: usize,
 23        k_ksk: usize,
 24        digits_ksk: usize,
 25        k_tsk: usize,
 26        digits_tsk: usize,
 27        rank: usize,
 28    ) -> usize
 29    where
 30        Module<B>:
 31            VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxBigAllocBytes + VecZnxNormalizeTmpBytes + VecZnxBigNormalizeTmpBytes,
 32    {
 33        let out_size: usize = k_out.div_ceil(basek);
 34        let ci_dft: usize = module.vec_znx_dft_alloc_bytes(rank + 1, out_size);
 35        let ks_internal: usize =
 36            GLWECiphertext::keyswitch_scratch_space(module, basek, k_out, k_in, k_ksk, digits_ksk, rank, rank);
 37        let expand: usize = GGSWCiphertext::expand_row_scratch_space(module, basek, k_out, k_tsk, digits_tsk, rank);
 38        ci_dft + (ks_internal | expand)
 39    }
 40
 41    #[allow(clippy::too_many_arguments)]
 42    pub fn automorphism_inplace_scratch_space<B: Backend>(
 43        module: &Module<B>,
 44        basek: usize,
 45        k_out: usize,
 46        k_ksk: usize,
 47        digits_ksk: usize,
 48        k_tsk: usize,
 49        digits_tsk: usize,
 50        rank: usize,
 51    ) -> usize
 52    where
 53        Module<B>:
 54            VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxBigAllocBytes + VecZnxNormalizeTmpBytes + VecZnxBigNormalizeTmpBytes,
 55    {
 56        GGSWCiphertext::automorphism_scratch_space(
 57            module, basek, k_out, k_out, k_ksk, digits_ksk, k_tsk, digits_tsk, rank,
 58        )
 59    }
 60}
 61
 62impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
 63    pub fn automorphism<DataLhs: DataRef, DataAk: DataRef, DataTsk: DataRef, B: Backend>(
 64        &mut self,
 65        module: &Module<B>,
 66        lhs: &GGSWCiphertext<DataLhs>,
 67        auto_key: &GGLWEAutomorphismKeyPrepared<DataAk, B>,
 68        tensor_key: &GGLWETensorKeyPrepared<DataTsk, B>,
 69        scratch: &mut Scratch<B>,
 70    ) where
 71        Module<B>: VecZnxDftAllocBytes
 72            + VmpApplyTmpBytes
 73            + VecZnxBigNormalizeTmpBytes
 74            + VmpApply<B>
 75            + VmpApplyAdd<B>
 76            + VecZnxDftFromVecZnx<B>
 77            + VecZnxDftToVecZnxBigConsume<B>
 78            + VecZnxBigAddSmallInplace<B>
 79            + VecZnxBigNormalize<B>
 80            + VecZnxAutomorphismInplace
 81            + VecZnxBigAllocBytes
 82            + VecZnxNormalizeTmpBytes
 83            + VecZnxDftCopy<B>
 84            + VecZnxDftAddInplace<B>
 85            + VecZnxDftToVecZnxBigTmpA<B>,
 86        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnxBig<B>,
 87    {
 88        #[cfg(debug_assertions)]
 89        {
 90            assert_eq!(self.n(), auto_key.n());
 91            assert_eq!(lhs.n(), auto_key.n());
 92
 93            assert_eq!(
 94                self.rank(),
 95                lhs.rank(),
 96                "ggsw_out rank: {} != ggsw_in rank: {}",
 97                self.rank(),
 98                lhs.rank()
 99            );
100            assert_eq!(
101                self.rank(),
102                auto_key.rank(),
103                "ggsw_in rank: {} != auto_key rank: {}",
104                self.rank(),
105                auto_key.rank()
106            );
107            assert_eq!(
108                self.rank(),
109                tensor_key.rank(),
110                "ggsw_in rank: {} != tensor_key rank: {}",
111                self.rank(),
112                tensor_key.rank()
113            );
114            assert!(
115                scratch.available()
116                    >= GGSWCiphertext::automorphism_scratch_space(
117                        module,
118                        self.basek(),
119                        self.k(),
120                        lhs.k(),
121                        auto_key.k(),
122                        auto_key.digits(),
123                        tensor_key.k(),
124                        tensor_key.digits(),
125                        self.rank(),
126                    )
127            )
128        };
129
130        self.automorphism_internal(module, lhs, auto_key, scratch);
131        self.expand_row(module, tensor_key, scratch);
132    }
133
134    pub fn automorphism_inplace<DataKsk: DataRef, DataTsk: DataRef, B: Backend>(
135        &mut self,
136        module: &Module<B>,
137        auto_key: &GGLWEAutomorphismKeyPrepared<DataKsk, B>,
138        tensor_key: &GGLWETensorKeyPrepared<DataTsk, B>,
139        scratch: &mut Scratch<B>,
140    ) where
141        Module<B>: VecZnxDftAllocBytes
142            + VmpApplyTmpBytes
143            + VecZnxBigNormalizeTmpBytes
144            + VmpApply<B>
145            + VmpApplyAdd<B>
146            + VecZnxDftFromVecZnx<B>
147            + VecZnxDftToVecZnxBigConsume<B>
148            + VecZnxBigAddSmallInplace<B>
149            + VecZnxBigNormalize<B>
150            + VecZnxAutomorphismInplace
151            + VecZnxBigAllocBytes
152            + VecZnxNormalizeTmpBytes
153            + VecZnxDftCopy<B>
154            + VecZnxDftAddInplace<B>
155            + VecZnxDftToVecZnxBigTmpA<B>,
156        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnxBig<B>,
157    {
158        unsafe {
159            let self_ptr: *mut GGSWCiphertext<DataSelf> = self as *mut GGSWCiphertext<DataSelf>;
160            self.automorphism(module, &*self_ptr, auto_key, tensor_key, scratch);
161        }
162    }
163
164    fn automorphism_internal<DataLhs: DataRef, DataAk: DataRef, B: Backend>(
165        &mut self,
166        module: &Module<B>,
167        lhs: &GGSWCiphertext<DataLhs>,
168        auto_key: &GGLWEAutomorphismKeyPrepared<DataAk, B>,
169        scratch: &mut Scratch<B>,
170    ) where
171        Module<B>: VecZnxDftAllocBytes
172            + VmpApplyTmpBytes
173            + VecZnxBigNormalizeTmpBytes
174            + VmpApply<B>
175            + VmpApplyAdd<B>
176            + VecZnxDftFromVecZnx<B>
177            + VecZnxDftToVecZnxBigConsume<B>
178            + VecZnxBigAddSmallInplace<B>
179            + VecZnxBigNormalize<B>
180            + VecZnxAutomorphismInplace,
181        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
182    {
183        // Keyswitch the j-th row of the col 0
184        (0..lhs.rows()).for_each(|row_i| {
185            // Key-switch column 0, i.e.
186            // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0pi^-1(s0) + a1pi^-1(s1) + a2pi^-1(s2)) + M[i], a0, a1, a2)
187            self.at_mut(row_i, 0)
188                .automorphism(module, &lhs.at(row_i, 0), auto_key, scratch);
189        });
190    }
191}