1use poulpy_hal::{
 2    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, ScalarZnxToRef, Scratch, ZnxView, ZnxViewMut},
 3    source::Source,
 4};
 5
 6use std::marker::PhantomData;
 7
 8use poulpy_core::{
 9    Distribution, GGSWEncryptSk, GetDistribution, ScratchTakeCore,
10    layouts::{GGSW, GGSWInfos, GLWEInfos, GLWESecretPreparedToRef, LWEInfos, LWESecret, LWESecretToRef},
11};
12
13use crate::bin_fhe::blind_rotation::{
14    BlindRotationKey, BlindRotationKeyEncryptSk, BlindRotationKeyFactory, BlindRotationKeyInfos, CGGI,
15};
16
17impl<D: DataRef> BlindRotationKeyFactory<CGGI> for BlindRotationKey<D, CGGI> {
18    fn blind_rotation_key_alloc<A>(infos: &A) -> BlindRotationKey<Vec<u8>, CGGI>
19    where
20        A: BlindRotationKeyInfos,
21    {
22        BlindRotationKey {
23            keys: (0..infos.n_lwe().as_usize()).map(|_| GGSW::alloc_from_infos(infos)).collect(),
24            dist: Distribution::NONE,
25            _phantom: PhantomData,
26        }
27    }
28}
29
30impl<BE: Backend> BlindRotationKeyEncryptSk<CGGI, BE> for Module<BE>
31where
32    Self: GGSWEncryptSk<BE>,
33    Scratch<BE>: ScratchTakeCore<BE>,
34{
35    fn blind_rotation_key_encrypt_sk_tmp_bytes<A: GGSWInfos>(&self, infos: &A) -> usize {
36        self.ggsw_encrypt_sk_tmp_bytes(infos)
37    }
38
39    fn blind_rotation_key_encrypt_sk<D, S0, S1>(
40        &self,
41        res: &mut BlindRotationKey<D, CGGI>,
42        sk_glwe: &S0,
43        sk_lwe: &S1,
44        source_xa: &mut Source,
45        source_xe: &mut Source,
46        scratch: &mut Scratch<BE>,
47    ) where
48        D: DataMut,
49        S0: GLWESecretPreparedToRef<BE> + GLWEInfos,
50        S1: LWESecretToRef + LWEInfos + GetDistribution,
51    {
52        assert_eq!(res.keys.len() as u32, sk_lwe.n());
53        assert!(sk_glwe.n() <= self.n() as u32);
54        assert_eq!(sk_glwe.rank(), res.rank());
55
56        match sk_lwe.dist() {
57            Distribution::BinaryBlock(_) | Distribution::BinaryFixed(_) | Distribution::BinaryProb(_) | Distribution::ZERO => {}
58            _ => {
59                panic!("invalid GLWESecret distribution: must be BinaryBlock, BinaryFixed or BinaryProb (or ZERO for debugging)")
60            }
61        }
62
63        {
64            let sk_lwe: &LWESecret<&[u8]> = &sk_lwe.to_ref();
65
66            res.dist = sk_lwe.dist();
67
68            let mut pt: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(sk_glwe.n().into(), 1);
69            let sk_ref: ScalarZnx<&[u8]> = sk_lwe.data().to_ref();
70
71            for (i, ggsw) in res.keys.iter_mut().enumerate() {
72                pt.at_mut(0, 0)[0] = sk_ref.at(0, 0)[i];
73                ggsw.encrypt_sk(self, &pt, sk_glwe, source_xa, source_xe, scratch);
74            }
75        }
76    }
77}