1use poulpy_hal::{
  2    api::{
  3        ScratchAvailable, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace, VecZnxBigBytesOf, VecZnxBigNormalize,
  4        VecZnxBigNormalizeTmpBytes, VecZnxDftAddInplace, VecZnxDftApply, VecZnxDftBytesOf, VecZnxDftCopy, VecZnxIdftApplyConsume,
  5        VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd,
  6        VmpApplyDftToDftTmpBytes,
  7    },
  8    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  9};
 10
 11use crate::layouts::{
 12    GGLWEInfos, GGSW, GGSWInfos, GLWE,
 13    prepared::{AutomorphismKeyPrepared, TensorKeyPrepared},
 14};
 15
 16impl GGSW<Vec<u8>> {
 17    pub fn automorphism_scratch_space<B: Backend, OUT, IN, KEY, TSK>(
 18        module: &Module<B>,
 19        out_infos: &OUT,
 20        in_infos: &IN,
 21        key_infos: &KEY,
 22        tsk_infos: &TSK,
 23    ) -> usize
 24    where
 25        OUT: GGSWInfos,
 26        IN: GGSWInfos,
 27        KEY: GGLWEInfos,
 28        TSK: GGLWEInfos,
 29        Module<B>:
 30            VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigBytesOf + VecZnxNormalizeTmpBytes + VecZnxBigNormalizeTmpBytes,
 31    {
 32        let out_size: usize = out_infos.size();
 33        let ci_dft: usize = module.bytes_of_vec_znx_dft((key_infos.rank_out() + 1).into(), out_size);
 34        let ks_internal: usize = GLWE::keyswitch_scratch_space(
 35            module,
 36            &out_infos.glwe_layout(),
 37            &in_infos.glwe_layout(),
 38            key_infos,
 39        );
 40        let expand: usize = GGSW::expand_row_scratch_space(module, out_infos, tsk_infos);
 41        ci_dft + (ks_internal | expand)
 42    }
 43
 44    pub fn automorphism_inplace_scratch_space<B: Backend, OUT, KEY, TSK>(
 45        module: &Module<B>,
 46        out_infos: &OUT,
 47        key_infos: &KEY,
 48        tsk_infos: &TSK,
 49    ) -> usize
 50    where
 51        OUT: GGSWInfos,
 52        KEY: GGLWEInfos,
 53        TSK: GGLWEInfos,
 54        Module<B>:
 55            VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigBytesOf + VecZnxNormalizeTmpBytes + VecZnxBigNormalizeTmpBytes,
 56    {
 57        GGSW::automorphism_scratch_space(module, out_infos, out_infos, key_infos, tsk_infos)
 58    }
 59}
 60
 61impl<DataSelf: DataMut> GGSW<DataSelf> {
 62    pub fn automorphism<DataLhs: DataRef, DataAk: DataRef, DataTsk: DataRef, B: Backend>(
 63        &mut self,
 64        module: &Module<B>,
 65        lhs: &GGSW<DataLhs>,
 66        auto_key: &AutomorphismKeyPrepared<DataAk, B>,
 67        tensor_key: &TensorKeyPrepared<DataTsk, B>,
 68        scratch: &mut Scratch<B>,
 69    ) where
 70        Module<B>: VecZnxDftBytesOf
 71            + VmpApplyDftToDftTmpBytes
 72            + VecZnxBigNormalizeTmpBytes
 73            + VmpApplyDftToDft<B>
 74            + VmpApplyDftToDftAdd<B>
 75            + VecZnxDftApply<B>
 76            + VecZnxIdftApplyConsume<B>
 77            + VecZnxBigAddSmallInplace<B>
 78            + VecZnxBigNormalize<B>
 79            + VecZnxAutomorphismInplace<B>
 80            + VecZnxBigBytesOf
 81            + VecZnxNormalizeTmpBytes
 82            + VecZnxDftCopy<B>
 83            + VecZnxDftAddInplace<B>
 84            + VecZnxIdftApplyTmpA<B>
 85            + VecZnxNormalize<B>,
 86        Scratch<B>: ScratchAvailable,
 87    {
 88        #[cfg(debug_assertions)]
 89        {
 90            use crate::layouts::{GLWEInfos, LWEInfos};
 91
 92            assert_eq!(self.n(), module.n() as u32);
 93            assert_eq!(lhs.n(), module.n() as u32);
 94            assert_eq!(auto_key.n(), module.n() as u32);
 95            assert_eq!(tensor_key.n(), module.n() as u32);
 96
 97            assert_eq!(
 98                self.rank(),
 99                lhs.rank(),
100                "ggsw_out rank: {} != ggsw_in rank: {}",
101                self.rank(),
102                lhs.rank()
103            );
104            assert_eq!(
105                self.rank(),
106                auto_key.rank_out(),
107                "ggsw_in rank: {} != auto_key rank: {}",
108                self.rank(),
109                auto_key.rank_out()
110            );
111            assert_eq!(
112                self.rank(),
113                tensor_key.rank_out(),
114                "ggsw_in rank: {} != tensor_key rank: {}",
115                self.rank(),
116                tensor_key.rank_out()
117            );
118            assert!(scratch.available() >= GGSW::automorphism_scratch_space(module, self, lhs, auto_key, tensor_key))
119        };
120
121        // Keyswitch the j-th row of the col 0
122        (0..lhs.dnum().into()).for_each(|row_i| {
123            // Key-switch column 0, i.e.
124            // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0pi^-1(s0) + a1pi^-1(s1) + a2pi^-1(s2)) + M[i], a0, a1, a2)
125            self.at_mut(row_i, 0)
126                .automorphism(module, &lhs.at(row_i, 0), auto_key, scratch);
127        });
128        self.expand_row(module, tensor_key, scratch);
129    }
130
131    pub fn automorphism_inplace<DataKsk: DataRef, DataTsk: DataRef, B: Backend>(
132        &mut self,
133        module: &Module<B>,
134        auto_key: &AutomorphismKeyPrepared<DataKsk, B>,
135        tensor_key: &TensorKeyPrepared<DataTsk, B>,
136        scratch: &mut Scratch<B>,
137    ) where
138        Module<B>: VecZnxDftBytesOf
139            + VmpApplyDftToDftTmpBytes
140            + VecZnxBigNormalizeTmpBytes
141            + VmpApplyDftToDft<B>
142            + VmpApplyDftToDftAdd<B>
143            + VecZnxDftApply<B>
144            + VecZnxIdftApplyConsume<B>
145            + VecZnxBigAddSmallInplace<B>
146            + VecZnxBigNormalize<B>
147            + VecZnxAutomorphismInplace<B>
148            + VecZnxBigBytesOf
149            + VecZnxNormalizeTmpBytes
150            + VecZnxDftCopy<B>
151            + VecZnxDftAddInplace<B>
152            + VecZnxIdftApplyTmpA<B>
153            + VecZnxNormalize<B>,
154        Scratch<B>: ScratchAvailable,
155    {
156        // Keyswitch the j-th row of the col 0
157        (0..self.dnum().into()).for_each(|row_i| {
158            // Key-switch column 0, i.e.
159            // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0pi^-1(s0) + a1pi^-1(s1) + a2pi^-1(s2)) + M[i], a0, a1, a2)
160            self.at_mut(row_i, 0)
161                .automorphism_inplace(module, auto_key, scratch);
162        });
163        self.expand_row(module, tensor_key, scratch);
164    }
165}