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<BE, CGGI> 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<DM, DR>(
37        &self,
38        res: &mut BlindRotationKeyPrepared<DM, CGGI, BE>,
39        other: &BlindRotationKey<DR, CGGI>,
40        scratch: &mut Scratch<BE>,
41    ) where
42        DM: DataMut,
43        DR: DataRef,
44    {
45        #[cfg(debug_assertions)]
46        {
47            assert_eq!(res.data.len(), other.keys.len());
48        }
49
50        let n: usize = other.n().as_usize();
51
52        for (a, b) in res.data.iter_mut().zip(other.keys.iter()) {
53            a.prepare(self, b, scratch);
54        }
55
56        res.dist = other.dist;
57
58        if let Distribution::BinaryBlock(_) = other.dist {
59            let mut x_pow_a: Vec<SvpPPol<Vec<u8>, BE>> = Vec::with_capacity(n << 1);
60            let mut buf: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
61            (0..n << 1).for_each(|i| {
62                let mut res: SvpPPol<Vec<u8>, BE> = self.svp_ppol_alloc(1);
63                set_xai_plus_y(self, i, 0, &mut res, &mut buf);
64                x_pow_a.push(res);
65            });
66            res.x_pow_a = Some(x_pow_a);
67        }
68    }
69}