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