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