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())
24                .map(|_| GGSW::alloc_from_infos(infos))
25                .collect(),
26            dist: Distribution::NONE,
27            _phantom: PhantomData,
28        }
29    }
30}
31
32impl<BE: Backend> BlindRotationKeyEncryptSk<CGGI, BE> for Module<BE>
33where
34    Self: GGSWEncryptSk<BE>,
35    Scratch<BE>: ScratchTakeCore<BE>,
36{
37    fn blind_rotation_key_encrypt_sk_tmp_bytes<A: GGSWInfos>(&self, infos: &A) -> usize {
38        self.ggsw_encrypt_sk_tmp_bytes(infos)
39    }
40
41    fn blind_rotation_key_encrypt_sk<D, S0, S1>(
42        &self,
43        res: &mut BlindRotationKey<D, CGGI>,
44        sk_glwe: &S0,
45        sk_lwe: &S1,
46        source_xa: &mut Source,
47        source_xe: &mut Source,
48        scratch: &mut Scratch<BE>,
49    ) where
50        D: DataMut,
51        S0: GLWESecretPreparedToRef<BE> + GLWEInfos,
52        S1: LWESecretToRef + LWEInfos + GetDistribution,
53    {
54        assert_eq!(res.keys.len() as u32, sk_lwe.n());
55        assert!(sk_glwe.n() <= self.n() as u32);
56        assert_eq!(sk_glwe.rank(), res.rank());
57
58        match sk_lwe.dist() {
59            Distribution::BinaryBlock(_) | Distribution::BinaryFixed(_) | Distribution::BinaryProb(_) | Distribution::ZERO => {}
60            _ => {
61                panic!("invalid GLWESecret distribution: must be BinaryBlock, BinaryFixed or BinaryProb (or ZERO for debugging)")
62            }
63        }
64
65        {
66            let sk_lwe: &LWESecret<&[u8]> = &sk_lwe.to_ref();
67
68            res.dist = sk_lwe.dist();
69
70            let mut pt: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(sk_glwe.n().into(), 1);
71            let sk_ref: ScalarZnx<&[u8]> = sk_lwe.data().to_ref();
72
73            for (i, ggsw) in res.keys.iter_mut().enumerate() {
74                pt.at_mut(0, 0)[0] = sk_ref.at(0, 0)[i];
75                ggsw.encrypt_sk(self, &pt, sk_glwe, source_xa, source_xe, scratch);
76            }
77        }
78    }
79}