1use poulpy_hal::{
  2    api::{SvpPPolAlloc, SvpPrepare, VmpPMatAlloc, VmpPrepare},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, ScalarZnx, Scratch, SvpPPol},
  4};
  5
  6use std::marker::PhantomData;
  7
  8use poulpy_core::{
  9    Distribution,
 10    layouts::{
 11        Base2K, Degree, Digits, GGSWInfos, GLWEInfos, LWEInfos, Rank, Rows, TorusPrecision,
 12        prepared::{GGSWCiphertextPrepared, Prepare, PrepareAlloc},
 13    },
 14};
 15
 16use crate::tfhe::blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyInfos, utils::set_xai_plus_y};
 17
 18pub trait BlindRotationKeyPreparedAlloc<B: Backend> {
 19    fn alloc<A>(module: &Module<B>, infos: &A) -> Self
 20    where
 21        A: BlindRotationKeyInfos;
 22}
 23
 24#[derive(PartialEq, Eq)]
 25pub struct BlindRotationKeyPrepared<D: Data, BRT: BlindRotationAlgo, B: Backend> {
 26    pub(crate) data: Vec<GGSWCiphertextPrepared<D, B>>,
 27    pub(crate) dist: Distribution,
 28    pub(crate) x_pow_a: Option<Vec<SvpPPol<Vec<u8>, B>>>,
 29    pub(crate) _phantom: PhantomData<BRT>,
 30}
 31
 32impl<D: Data, BRT: BlindRotationAlgo, B: Backend> BlindRotationKeyInfos for BlindRotationKeyPrepared<D, BRT, B> {
 33    fn n_glwe(&self) -> Degree {
 34        self.n()
 35    }
 36
 37    fn n_lwe(&self) -> Degree {
 38        Degree(self.data.len() as u32)
 39    }
 40}
 41
 42impl<D: Data, BRT: BlindRotationAlgo, B: Backend> LWEInfos for BlindRotationKeyPrepared<D, BRT, B> {
 43    fn base2k(&self) -> Base2K {
 44        self.data[0].base2k()
 45    }
 46
 47    fn k(&self) -> TorusPrecision {
 48        self.data[0].k()
 49    }
 50
 51    fn n(&self) -> Degree {
 52        self.data[0].n()
 53    }
 54
 55    fn size(&self) -> usize {
 56        self.data[0].size()
 57    }
 58}
 59
 60impl<D: Data, BRT: BlindRotationAlgo, B: Backend> GLWEInfos for BlindRotationKeyPrepared<D, BRT, B> {
 61    fn rank(&self) -> Rank {
 62        self.data[0].rank()
 63    }
 64}
 65impl<D: Data, BRT: BlindRotationAlgo, B: Backend> GGSWInfos for BlindRotationKeyPrepared<D, BRT, B> {
 66    fn digits(&self) -> poulpy_core::layouts::Digits {
 67        Digits(1)
 68    }
 69
 70    fn rows(&self) -> Rows {
 71        self.data[0].rows()
 72    }
 73}
 74
 75impl<D: Data, BRT: BlindRotationAlgo, B: Backend> BlindRotationKeyPrepared<D, BRT, B> {
 76    pub fn block_size(&self) -> usize {
 77        match self.dist {
 78            Distribution::BinaryBlock(value) => value,
 79            _ => 1,
 80        }
 81    }
 82}
 83
 84impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>
 85    for BlindRotationKey<D, BRA>
 86where
 87    BlindRotationKeyPrepared<Vec<u8>, BRA, B>: BlindRotationKeyPreparedAlloc<B>,
 88    BlindRotationKeyPrepared<Vec<u8>, BRA, B>: Prepare<B, BlindRotationKey<D, BRA>>,
 89{
 90    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> BlindRotationKeyPrepared<Vec<u8>, BRA, B> {
 91        let mut brk: BlindRotationKeyPrepared<Vec<u8>, BRA, B> = BlindRotationKeyPrepared::alloc(module, self);
 92        brk.prepare(module, self, scratch);
 93        brk
 94    }
 95}
 96
 97impl<DM: DataMut, DR: DataRef, BRA: BlindRotationAlgo, B: Backend> Prepare<B, BlindRotationKey<DR, BRA>>
 98    for BlindRotationKeyPrepared<DM, BRA, B>
 99where
100    Module<B>: VmpPMatAlloc<B> + VmpPrepare<B> + SvpPPolAlloc<B> + SvpPrepare<B>,
101{
102    fn prepare(&mut self, module: &Module<B>, other: &BlindRotationKey<DR, BRA>, scratch: &mut Scratch<B>) {
103        #[cfg(debug_assertions)]
104        {
105            assert_eq!(self.data.len(), other.keys.len());
106        }
107
108        let n: usize = other.n().as_usize();
109
110        self.data
111            .iter_mut()
112            .zip(other.keys.iter())
113            .for_each(|(ggsw_prepared, other)| {
114                ggsw_prepared.prepare(module, other, scratch);
115            });
116
117        self.dist = other.dist;
118
119        if let Distribution::BinaryBlock(_) = other.dist {
120            let mut x_pow_a: Vec<SvpPPol<Vec<u8>, B>> = Vec::with_capacity(n << 1);
121            let mut buf: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
122            (0..n << 1).for_each(|i| {
123                let mut res: SvpPPol<Vec<u8>, B> = module.svp_ppol_alloc(1);
124                set_xai_plus_y(module, i, 0, &mut res, &mut buf);
125                x_pow_a.push(res);
126            });
127            self.x_pow_a = Some(x_pow_a);
128        }
129    }
130}