1use poulpy_hal::layouts::{Backend, Data, DataMut, DataRef, Scratch, SvpPPol};
  2
  3use std::marker::PhantomData;
  4
  5use poulpy_core::{
  6    Distribution,
  7    layouts::{Base2K, Degree, Dnum, Dsize, GGSWInfos, GLWEInfos, LWEInfos, Rank, TorusPrecision, prepared::GGSWPrepared},
  8};
  9
 10use crate::tfhe::blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyInfos};
 11
 12pub trait BlindRotationKeyPreparedFactory<BRA: BlindRotationAlgo, BE: Backend> {
 13    fn blind_rotation_key_prepared_alloc<A>(&self, infos: &A) -> BlindRotationKeyPrepared<Vec<u8>, BRA, BE>
 14    where
 15        A: BlindRotationKeyInfos;
 16
 17    fn blind_rotation_key_prepare<DM, DR>(
 18        &self,
 19        res: &mut BlindRotationKeyPrepared<DM, BRA, BE>,
 20        other: &BlindRotationKey<DR, BRA>,
 21        scratch: &mut Scratch<BE>,
 22    ) where
 23        DM: DataMut,
 24        DR: DataRef;
 25}
 26
 27impl<BE: Backend, BRA: BlindRotationAlgo> BlindRotationKeyPrepared<Vec<u8>, BRA, BE> {
 28    pub fn alloc<A, M>(module: &M, infos: &A) -> Self
 29    where
 30        A: BlindRotationKeyInfos,
 31        M: BlindRotationKeyPreparedFactory<BRA, BE>,
 32    {
 33        module.blind_rotation_key_prepared_alloc(infos)
 34    }
 35}
 36
 37impl<D: DataMut, BRA: BlindRotationAlgo, BE: Backend> BlindRotationKeyPrepared<D, BRA, BE> {
 38    pub fn prepare<DR: DataRef, M>(&mut self, module: &M, other: &BlindRotationKey<DR, BRA>, scratch: &mut Scratch<BE>)
 39    where
 40        M: BlindRotationKeyPreparedFactory<BRA, BE>,
 41    {
 42        module.blind_rotation_key_prepare(self, other, scratch);
 43    }
 44}
 45
 46#[derive(PartialEq, Eq)]
 47pub struct BlindRotationKeyPrepared<D: Data, BRT: BlindRotationAlgo, B: Backend> {
 48    pub(crate) data: Vec<GGSWPrepared<D, B>>,
 49    pub(crate) dist: Distribution,
 50    pub(crate) x_pow_a: Option<Vec<SvpPPol<Vec<u8>, B>>>,
 51    pub(crate) _phantom: PhantomData<BRT>,
 52}
 53
 54impl<D: Data, BRT: BlindRotationAlgo, B: Backend> BlindRotationKeyInfos for BlindRotationKeyPrepared<D, BRT, B> {
 55    fn n_glwe(&self) -> Degree {
 56        self.n()
 57    }
 58
 59    fn n_lwe(&self) -> Degree {
 60        Degree(self.data.len() as u32)
 61    }
 62}
 63
 64impl<D: Data, BRT: BlindRotationAlgo, B: Backend> LWEInfos for BlindRotationKeyPrepared<D, BRT, B> {
 65    fn base2k(&self) -> Base2K {
 66        self.data[0].base2k()
 67    }
 68
 69    fn k(&self) -> TorusPrecision {
 70        self.data[0].k()
 71    }
 72
 73    fn n(&self) -> Degree {
 74        self.data[0].n()
 75    }
 76
 77    fn size(&self) -> usize {
 78        self.data[0].size()
 79    }
 80}
 81
 82impl<D: Data, BRT: BlindRotationAlgo, B: Backend> GLWEInfos for BlindRotationKeyPrepared<D, BRT, B> {
 83    fn rank(&self) -> Rank {
 84        self.data[0].rank()
 85    }
 86}
 87impl<D: Data, BRT: BlindRotationAlgo, B: Backend> GGSWInfos for BlindRotationKeyPrepared<D, BRT, B> {
 88    fn dsize(&self) -> poulpy_core::layouts::Dsize {
 89        Dsize(1)
 90    }
 91
 92    fn dnum(&self) -> Dnum {
 93        self.data[0].dnum()
 94    }
 95}
 96
 97impl<D: Data, BRT: BlindRotationAlgo, B: Backend> BlindRotationKeyPrepared<D, BRT, B> {
 98    pub fn block_size(&self) -> usize {
 99        match self.dist {
100            Distribution::BinaryBlock(value) => value,
101            _ => 1,
102        }
103    }
104}