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