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