1use poulpy_hal::{
  2    api::{
  3        ModuleN, ScratchAvailable, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolBytesOf, SvpPrepare, VecZnxAddInplace, VecZnxAddNormal,
  4        VecZnxAddScalarInplace, VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf, VecZnxFillUniform, VecZnxIdftApplyConsume,
  5        VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace, VecZnxSwitchRing,
  6        ScratchTakeBasic,
  7    },
  8    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, Scratch},
  9    source::Source,
 10};
 11
 12use crate::{
 13    ScratchTakeCore,
 14    layouts::{
 15        GGLWE, GGLWEInfos, GLWEInfos, GLWESecret, GLWESwitchingKey, LWEInfos, RingDegree, prepared::GLWESecretPrepared,
 16    },
 17    encryption::gglwe_ct::GGLWEEncryptSk,
 18};
 19
 20impl GLWESwitchingKey<Vec<u8>> {
 21    pub fn encrypt_sk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
 22    where
 23        A: GGLWEInfos,
 24        M: GLWESwitchingKeyEncryptSk<BE>,
 25    {
 26        module.glwe_switching_key_encrypt_sk_tmp_bytes(infos)
 27    }
 28
 29
 30    pub fn encrypt_pk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
 31    where
 32        A: GGLWEInfos,
 33        M: GLWESwitchingKeyEncryptSk<BE>,
 34    {
 35        module.glwe_switching_key_encrypt_pk_tmp_bytes(infos)
 36    }
 37
 38    // pub fn encrypt_sk_tmp_bytes<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
 39    // where
 40    //     A: GGLWEInfos,
 41    //     Module<B>: ModuleN + SvpPPolBytesOf + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes + SvpPPolAlloc<B>,
 42    // {
 43    //     (GGLWE::encrypt_sk_tmp_bytes(module, infos) | ScalarZnx::bytes_of(module.n(), 1))
 44    //         + ScalarZnx::bytes_of(module.n(), infos.rank_in().into())
 45    //         + GLWESecretPrepared::bytes_of_from_infos(module, &infos.glwe_layout())
 46    // }
 47
 48    // pub fn encrypt_pk_tmp_bytes<B: Backend, A>(module: &Module<B>, _infos: &A) -> usize
 49    // where
 50    //     A: GGLWEInfos,
 51    // {
 52    //     GGLWE::encrypt_pk_tmp_bytes(module, _infos)
 53    // }
 54}
 55
 56impl<DataSelf: DataMut> GLWESwitchingKey<DataSelf> {
 57    #[allow(clippy::too_many_arguments)]
 58    pub fn encrypt_sk<DataSkIn: DataRef, DataSkOut: DataRef, B: Backend>(
 59        &mut self,
 60        module: &Module<B>,
 61        sk_in: &GLWESecret<DataSkIn>,
 62        sk_out: &GLWESecret<DataSkOut>,
 63        source_xa: &mut Source,
 64        source_xe: &mut Source,
 65        scratch: &mut Scratch<B>,
 66    ) where
 67        Module<B>: ModuleN
 68            + VecZnxAddScalarInplace
 69            + VecZnxDftBytesOf
 70            + VecZnxBigNormalize<B>
 71            + VecZnxDftApply<B>
 72            + SvpApplyDftToDftInplace<B>
 73            + VecZnxIdftApplyConsume<B>
 74            + VecZnxNormalizeTmpBytes
 75            + VecZnxFillUniform
 76            + VecZnxSubInplace
 77            + VecZnxAddInplace
 78            + VecZnxNormalizeInplace<B>
 79            + VecZnxAddNormal
 80            + VecZnxNormalize<B>
 81            + VecZnxSub
 82            + SvpPrepare<B>
 83            + VecZnxSwitchRing
 84            + SvpPPolBytesOf
 85            + SvpPPolAlloc<B>,
 86        Scratch<B>: ScratchAvailable + ScratchTakeBasic + ScratchTakeCore<B>,
 87    {
 88        #[cfg(debug_assertions)]
 89        {
 90            assert!(sk_in.n().0 <= module.n() as u32);
 91            assert!(sk_out.n().0 <= module.n() as u32);
 92            assert!(
 93                scratch.available() >= GLWESwitchingKey::encrypt_sk_tmp_bytes(module, self),
 94                "scratch.available()={} < GLWESwitchingKey::encrypt_sk_tmp_bytes={}",
 95                scratch.available(),
 96                GLWESwitchingKey::encrypt_sk_tmp_bytes(module, self)
 97            )
 98        }
 99
100        let n: usize = sk_in.n().max(sk_out.n()).into();
101
102        let (mut sk_in_tmp, scratch_1) = scratch.take_scalar_znx(module, sk_in.rank().into());
103        (0..sk_in.rank().into()).for_each(|i| {
104            module.vec_znx_switch_ring(
105                &mut sk_in_tmp.as_vec_znx_mut(),
106                i,
107                &sk_in.data.as_vec_znx(),
108                i,
109            );
110        });
111
112        let (mut sk_out_tmp, scratch_2) = scratch_1.take_glwe_secret_prepared(module, sk_out.rank());
113        {
114            let (mut tmp, _) = scratch_2.take_scalar_znx(module, 1);
115            (0..sk_out.rank().into()).for_each(|i| {
116                module.vec_znx_switch_ring(&mut tmp.as_vec_znx_mut(), 0, &sk_out.data.as_vec_znx(), i);
117                module.svp_prepare(&mut sk_out_tmp.data, i, &tmp, 0);
118            });
119        }
120
121        self.key.encrypt_sk(
122            module,
123            &sk_in_tmp,
124            &sk_out_tmp,
125            source_xa,
126            source_xe,
127            scratch_2,
128        );
129        self.sk_in_n = sk_in.n().into();
130        self.sk_out_n = sk_out.n().into();
131    }
132}
133
134
135pub trait GLWESwitchingKeyEncryptSk<BE: Backend> 
136where
137    Self: Sized
138        + ModuleN
139        + SvpPPolBytesOf
140        + VecZnxDftBytesOf
141        + VecZnxNormalizeTmpBytes
142        + SvpPPolAlloc<BE>
143        + GGLWEEncryptSk<BE>,
144{
145    fn glwe_switching_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
146    where
147        A: GGLWEInfos;
148
149    fn glwe_switching_key_encrypt_pk_tmp_bytes<A>(&self, infos: &A) -> usize
150    where
151        A: GGLWEInfos;
152}
153
154impl<BE: Backend> GLWESwitchingKeyEncryptSk<BE> for Module<BE> where
155    Self: ModuleN
156        + SvpPPolBytesOf
157        + VecZnxDftBytesOf
158        + VecZnxNormalizeTmpBytes
159        + SvpPPolAlloc<BE>
160        + GGLWEEncryptSk<BE>,
161{
162
163    fn glwe_switching_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
164    where
165        A: GGLWEInfos,
166    {
167        (GGLWE::encrypt_sk_tmp_bytes(self, infos) | ScalarZnx::bytes_of(self.n(), 1))
168            + ScalarZnx::bytes_of(self.n(), infos.rank_in().into())
169            + GLWESecretPrepared::bytes_of_from_infos(self, &infos.glwe_layout())
170    }
171
172    fn glwe_switching_key_encrypt_pk_tmp_bytes<A>(&self, infos: &A) -> usize
173    where
174        A: GGLWEInfos,
175    {
176        GGLWE::encrypt_pk_tmp_bytes(self, infos)
177    }
178
179}