1use poulpy_hal::{
  2    api::ScratchAvailable,
  3    layouts::{Backend, DataMut, Module, Scratch},
  4};
  5
  6use crate::{
  7    GGSWExpandRows, ScratchTakeCore,
  8    automorphism::glwe_ct::GLWEAutomorphism,
  9    layouts::{
 10        GGLWEInfos, GGSW, GGSWInfos, GGSWToMut, GGSWToRef, GLWEInfos, LWEInfos,
 11        prepared::{AutomorphismKeyPrepared, AutomorphismKeyPreparedToRef, TensorKeyPrepared, TensorKeyPreparedToRef},
 12    },
 13};
 14
 15impl GGSW<Vec<u8>> {
 16    pub fn automorphism_tmp_bytes<R, A, K, T, M, BE: Backend>(
 17        module: &M,
 18        res_infos: &R,
 19        a_infos: &A,
 20        key_infos: &K,
 21        tsk_infos: &T,
 22    ) -> usize
 23    where
 24        R: GGSWInfos,
 25        A: GGSWInfos,
 26        K: GGLWEInfos,
 27        T: GGLWEInfos,
 28        M: GGSWAutomorphism<BE>,
 29    {
 30        module.ggsw_automorphism_tmp_bytes(res_infos, a_infos, key_infos, tsk_infos)
 31    }
 32}
 33
 34impl<D: DataMut> GGSW<D> {
 35    pub fn automorphism<A, K, T, M, BE: Backend>(&mut self, module: &M, a: &A, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
 36    where
 37        A: GGSWToRef,
 38        K: AutomorphismKeyPreparedToRef<BE>,
 39        T: TensorKeyPreparedToRef<BE>,
 40        Scratch<BE>: ScratchTakeCore<BE>,
 41        M: GGSWAutomorphism<BE>,
 42    {
 43        module.ggsw_automorphism(self, a, key, tsk, scratch);
 44    }
 45
 46    pub fn automorphism_inplace<K, T, M, BE: Backend>(&mut self, module: &M, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
 47    where
 48        K: AutomorphismKeyPreparedToRef<BE>,
 49        T: TensorKeyPreparedToRef<BE>,
 50        Scratch<BE>: ScratchTakeCore<BE>,
 51        M: GGSWAutomorphism<BE>,
 52    {
 53        module.ggsw_automorphism_inplace(self, key, tsk, scratch);
 54    }
 55}
 56
 57impl<BE: Backend> GGSWAutomorphism<BE> for Module<BE> where Self: GLWEAutomorphism<BE> + GGSWExpandRows<BE> {}
 58
 59pub trait GGSWAutomorphism<BE: Backend>
 60where
 61    Self: GLWEAutomorphism<BE> + GGSWExpandRows<BE>,
 62{
 63    fn ggsw_automorphism_tmp_bytes<R, A, K, T>(&self, res_infos: &R, a_infos: &A, key_infos: &K, tsk_infos: &T) -> usize
 64    where
 65        R: GGSWInfos,
 66        A: GGSWInfos,
 67        K: GGLWEInfos,
 68        T: GGLWEInfos,
 69    {
 70        let out_size: usize = res_infos.size();
 71        let ci_dft: usize = self.bytes_of_vec_znx_dft((key_infos.rank_out() + 1).into(), out_size);
 72        let ks_internal: usize = self.glwe_automorphism_tmp_bytes(res_infos, a_infos, key_infos);
 73        let expand: usize = self.ggsw_expand_rows_tmp_bytes(res_infos, tsk_infos);
 74        ci_dft + (ks_internal.max(expand))
 75    }
 76
 77    fn ggsw_automorphism<R, A, K, T>(&self, res: &mut R, a: &A, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
 78    where
 79        R: GGSWToMut,
 80        A: GGSWToRef,
 81        K: AutomorphismKeyPreparedToRef<BE>,
 82        T: TensorKeyPreparedToRef<BE>,
 83        Scratch<BE>: ScratchTakeCore<BE>,
 84    {
 85        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
 86        let a: &GGSW<&[u8]> = &a.to_ref();
 87        let key: &AutomorphismKeyPrepared<&[u8], BE> = &key.to_ref();
 88        let tsk: &TensorKeyPrepared<&[u8], BE> = &tsk.to_ref();
 89
 90        assert_eq!(res.ggsw_layout(), a.ggsw_layout());
 91        assert_eq!(res.glwe_layout(), a.glwe_layout());
 92        assert_eq!(res.lwe_layout(), a.lwe_layout());
 93        assert!(scratch.available() >= self.ggsw_automorphism_tmp_bytes(res, a, key, tsk));
 94
 95        // Keyswitch the j-th row of the col 0
 96        for row in 0..res.dnum().as_usize() {
 97            // Key-switch column 0, i.e.
 98            // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0pi^-1(s0) + a1pi^-1(s1) + a2pi^-1(s2)) + M[i], a0, a1, a2)
 99            self.glwe_automorphism(&mut res.at_mut(row, 0), &a.at(row, 0), key, scratch);
100        }
101
102        self.ggsw_expand_row(res, tsk, scratch);
103    }
104
105    fn ggsw_automorphism_inplace<R, K, T>(&self, res: &mut R, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
106    where
107        R: GGSWToMut,
108        K: AutomorphismKeyPreparedToRef<BE>,
109        T: TensorKeyPreparedToRef<BE>,
110        Scratch<BE>: ScratchTakeCore<BE>,
111    {
112        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
113        let key: &AutomorphismKeyPrepared<&[u8], BE> = &key.to_ref();
114        let tsk: &TensorKeyPrepared<&[u8], BE> = &tsk.to_ref();
115
116        // Keyswitch the j-th row of the col 0
117        for row in 0..res.dnum().as_usize() {
118            // Key-switch column 0, i.e.
119            // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0pi^-1(s0) + a1pi^-1(s1) + a2pi^-1(s2)) + M[i], a0, a1, a2)
120            self.glwe_automorphism_inplace(&mut res.at_mut(row, 0), key, scratch);
121        }
122
123        self.ggsw_expand_row(res, tsk, scratch);
124    }
125}
126
127impl<DataSelf: DataMut> GGSW<DataSelf> {}