1use backend::hal::{
  2    api::{ScratchAvailable, TakeVecZnxBig, TakeVecZnxDft, VecZnxAutomorphismInplace, VecZnxNormalizeTmpBytes},
  3    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  4};
  5
  6use crate::{
  7    AutomorphismKeyExec, GGSWCiphertext, GGSWKeySwitchFamily, GLWECiphertext, GLWEKeyswitchFamily, GLWETensorKeyExec, Infos,
  8};
  9
 10impl GGSWCiphertext<Vec<u8>> {
 11    pub fn automorphism_scratch_space<B: Backend>(
 12        module: &Module<B>,
 13        basek: usize,
 14        k_out: usize,
 15        k_in: usize,
 16        k_ksk: usize,
 17        digits_ksk: usize,
 18        k_tsk: usize,
 19        digits_tsk: usize,
 20        rank: usize,
 21    ) -> usize
 22    where
 23        Module<B>: GLWEKeyswitchFamily<B> + GGSWKeySwitchFamily<B> + VecZnxNormalizeTmpBytes,
 24    {
 25        let out_size: usize = k_out.div_ceil(basek);
 26        let ci_dft: usize = module.vec_znx_dft_alloc_bytes(rank + 1, out_size);
 27        let ks_internal: usize =
 28            GLWECiphertext::keyswitch_scratch_space(module, basek, k_out, k_in, k_ksk, digits_ksk, rank, rank);
 29        let expand: usize = GGSWCiphertext::expand_row_scratch_space(module, basek, k_out, k_tsk, digits_tsk, rank);
 30        ci_dft + (ks_internal | expand)
 31    }
 32
 33    pub fn automorphism_inplace_scratch_space<B: Backend>(
 34        module: &Module<B>,
 35        basek: usize,
 36        k_out: usize,
 37        k_ksk: usize,
 38        digits_ksk: usize,
 39        k_tsk: usize,
 40        digits_tsk: usize,
 41        rank: usize,
 42    ) -> usize
 43    where
 44        Module<B>: GLWEKeyswitchFamily<B> + GGSWKeySwitchFamily<B> + VecZnxNormalizeTmpBytes,
 45    {
 46        GGSWCiphertext::automorphism_scratch_space(
 47            module, basek, k_out, k_out, k_ksk, digits_ksk, k_tsk, digits_tsk, rank,
 48        )
 49    }
 50}
 51
 52impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
 53    pub fn automorphism<DataLhs: DataRef, DataAk: DataRef, DataTsk: DataRef, B: Backend>(
 54        &mut self,
 55        module: &Module<B>,
 56        lhs: &GGSWCiphertext<DataLhs>,
 57        auto_key: &AutomorphismKeyExec<DataAk, B>,
 58        tensor_key: &GLWETensorKeyExec<DataTsk, B>,
 59        scratch: &mut Scratch<B>,
 60    ) where
 61        Module<B>: GLWEKeyswitchFamily<B> + GGSWKeySwitchFamily<B> + VecZnxAutomorphismInplace + VecZnxNormalizeTmpBytes,
 62        Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B>,
 63    {
 64        #[cfg(debug_assertions)]
 65        {
 66            use crate::Infos;
 67
 68            assert_eq!(
 69                self.rank(),
 70                lhs.rank(),
 71                "ggsw_out rank: {} != ggsw_in rank: {}",
 72                self.rank(),
 73                lhs.rank()
 74            );
 75            assert_eq!(
 76                self.rank(),
 77                auto_key.rank(),
 78                "ggsw_in rank: {} != auto_key rank: {}",
 79                self.rank(),
 80                auto_key.rank()
 81            );
 82            assert_eq!(
 83                self.rank(),
 84                tensor_key.rank(),
 85                "ggsw_in rank: {} != tensor_key rank: {}",
 86                self.rank(),
 87                tensor_key.rank()
 88            );
 89            assert!(
 90                scratch.available()
 91                    >= GGSWCiphertext::automorphism_scratch_space(
 92                        module,
 93                        self.basek(),
 94                        self.k(),
 95                        lhs.k(),
 96                        auto_key.k(),
 97                        auto_key.digits(),
 98                        tensor_key.k(),
 99                        tensor_key.digits(),
100                        self.rank(),
101                    )
102            )
103        };
104
105        let rank: usize = self.rank();
106        let cols: usize = rank + 1;
107
108        // Keyswitch the j-th row of the col 0
109        (0..lhs.rows()).for_each(|row_i| {
110            // Key-switch column 0, i.e.
111            // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0pi^-1(s0) + a1pi^-1(s1) + a2pi^-1(s2)) + M[i], a0, a1, a2)
112            self.at_mut(row_i, 0)
113                .automorphism(module, &lhs.at(row_i, 0), auto_key, scratch);
114
115            // Isolates DFT(AUTO(a[i]))
116            let (mut ci_dft, scratch1) = scratch.take_vec_znx_dft(module, cols, self.size());
117            (0..cols).for_each(|i| {
118                module.vec_znx_dft_from_vec_znx(1, 0, &mut ci_dft, i, &self.at(row_i, 0).data, i);
119            });
120
121            // Generates
122            //
123            // col 1: (-(b0s0 + b1s1 + b2s2)    , b0 + pi(M[i]), b1           , b2           )
124            // col 2: (-(c0s0 + c1s1 + c2s2)    , c0           , c1 + pi(M[i]), c2           )
125            // col 3: (-(d0s0 + d1s1 + d2s2)    , d0           , d1           , d2 + pi(M[i]))
126            (1..cols).for_each(|col_j| {
127                self.expand_row(module, row_i, col_j, &ci_dft, tensor_key, scratch1);
128            });
129        })
130    }
131
132    pub fn automorphism_inplace<DataKsk: DataRef, DataTsk: DataRef, B: Backend>(
133        &mut self,
134        module: &Module<B>,
135        auto_key: &AutomorphismKeyExec<DataKsk, B>,
136        tensor_key: &GLWETensorKeyExec<DataTsk, B>,
137        scratch: &mut Scratch<B>,
138    ) where
139        Module<B>: GLWEKeyswitchFamily<B> + GGSWKeySwitchFamily<B> + VecZnxAutomorphismInplace + VecZnxNormalizeTmpBytes,
140        Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B>,
141    {
142        unsafe {
143            let self_ptr: *mut GGSWCiphertext<DataSelf> = self as *mut GGSWCiphertext<DataSelf>;
144            self.automorphism(module, &*self_ptr, auto_key, tensor_key, scratch);
145        }
146    }
147}