1use poulpy_hal::{
 2    api::{SvpPPolAlloc, SvpPrepare},
 3    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, Scratch, SvpPPol},
 4};
 5
 6use std::marker::PhantomData;
 7
 8use poulpy_core::{
 9    Distribution,
10    layouts::{GGSWPreparedFactory, LWEInfos, prepared::GGSWPrepared},
11};
12
13use crate::tfhe::blind_rotation::{
14    BlindRotationKey, BlindRotationKeyInfos, BlindRotationKeyPrepared, BlindRotationKeyPreparedFactory, CGGI,
15    utils::set_xai_plus_y,
16};
17
18impl<BE: Backend> BlindRotationKeyPreparedFactory<CGGI, BE> for Module<BE>
19where
20    Self: GGSWPreparedFactory<BE> + SvpPPolAlloc<BE> + SvpPrepare<BE>,
21{
22    fn blind_rotation_key_prepared_alloc<A>(&self, infos: &A) -> BlindRotationKeyPrepared<Vec<u8>, CGGI, BE>
23    where
24        A: BlindRotationKeyInfos,
25    {
26        BlindRotationKeyPrepared {
27            data: (0..infos.n_lwe().as_usize())
28                .map(|_| GGSWPrepared::alloc_from_infos(self, infos))
29                .collect(),
30            dist: Distribution::NONE,
31            x_pow_a: None,
32            _phantom: PhantomData,
33        }
34    }
35
36    fn blind_rotation_key_prepare_tmp_bytes<A>(&self, infos: &A) -> usize
37    where
38        A: BlindRotationKeyInfos,
39    {
40        self.ggsw_prepare_tmp_bytes(infos)
41    }
42
43    fn prepare_blind_rotation_key<DM, DR>(
44        &self,
45        res: &mut BlindRotationKeyPrepared<DM, CGGI, BE>,
46        other: &BlindRotationKey<DR, CGGI>,
47        scratch: &mut Scratch<BE>,
48    ) where
49        DM: DataMut,
50        DR: DataRef,
51    {
52        #[cfg(debug_assertions)]
53        {
54            assert_eq!(res.data.len(), other.keys.len());
55        }
56
57        let n: usize = other.n().as_usize();
58
59        for (a, b) in res.data.iter_mut().zip(other.keys.iter()) {
60            a.prepare(self, b, scratch);
61        }
62
63        res.dist = other.dist;
64
65        if let Distribution::BinaryBlock(_) = other.dist {
66            let mut x_pow_a: Vec<SvpPPol<Vec<u8>, BE>> = Vec::with_capacity(n << 1);
67            let mut buf: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
68            (0..n << 1).for_each(|i| {
69                let mut res: SvpPPol<Vec<u8>, BE> = self.svp_ppol_alloc(1);
70                set_xai_plus_y(self, i, 0, &mut res, &mut buf);
71                x_pow_a.push(res);
72            });
73            res.x_pow_a = Some(x_pow_a);
74        }
75    }
76}