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