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