1use poulpy_hal::{
  2    layouts::{Backend, Data, DataMut, DataRef, FillUniform, Module, ReaderFrom, Reset, Scratch, WriterTo},
  3    source::Source,
  4};
  5
  6use std::{fmt, marker::PhantomData};
  7
  8use poulpy_core::{
  9    Distribution,
 10    layouts::{GGSWCiphertext, Infos, LWESecret, prepared::GLWESecretPrepared},
 11};
 12
 13use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 14
 15use crate::tfhe::blind_rotation::BlindRotationAlgo;
 16
 17pub trait BlindRotationKeyAlloc {
 18    fn alloc(n_gglwe: usize, n_lwe: usize, basek: usize, k: usize, rows: usize, rank: usize) -> Self;
 19}
 20
 21pub trait BlindRotationKeyEncryptSk<B: Backend> {
 22    #[allow(clippy::too_many_arguments)]
 23    fn encrypt_sk<DataSkGLWE, DataSkLWE>(
 24        &mut self,
 25        module: &Module<B>,
 26        sk_glwe: &GLWESecretPrepared<DataSkGLWE, B>,
 27        sk_lwe: &LWESecret<DataSkLWE>,
 28        source_xa: &mut Source,
 29        source_xe: &mut Source,
 30        scratch: &mut Scratch<B>,
 31    ) where
 32        DataSkGLWE: DataRef,
 33        DataSkLWE: DataRef;
 34}
 35
 36#[derive(Clone)]
 37pub struct BlindRotationKey<D: Data, BRT: BlindRotationAlgo> {
 38    pub(crate) keys: Vec<GGSWCiphertext<D>>,
 39    pub(crate) dist: Distribution,
 40    pub(crate) _phantom: PhantomData<BRT>,
 41}
 42
 43impl<D: DataRef, BRT: BlindRotationAlgo> fmt::Debug for BlindRotationKey<D, BRT> {
 44    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 45        write!(f, "{}", self)
 46    }
 47}
 48
 49impl<D: Data, BRT: BlindRotationAlgo> PartialEq for BlindRotationKey<D, BRT> {
 50    fn eq(&self, other: &Self) -> bool {
 51        if self.keys.len() != other.keys.len() {
 52            return false;
 53        }
 54        for (a, b) in self.keys.iter().zip(other.keys.iter()) {
 55            if a != b {
 56                return false;
 57            }
 58        }
 59
 60        self.dist == other.dist && self._phantom == other._phantom
 61    }
 62}
 63
 64impl<D: Data, BRT: BlindRotationAlgo> Eq for BlindRotationKey<D, BRT> {}
 65
 66impl<D: DataRef, BRT: BlindRotationAlgo> fmt::Display for BlindRotationKey<D, BRT> {
 67    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 68        for (i, key) in self.keys.iter().enumerate() {
 69            write!(f, "key[{}]: {}", i, key)?;
 70        }
 71        writeln!(f, "{:?}", self.dist)
 72    }
 73}
 74
 75impl<D: DataMut, BRT: BlindRotationAlgo> Reset for BlindRotationKey<D, BRT> {
 76    fn reset(&mut self) {
 77        self.keys.iter_mut().for_each(|key| key.reset());
 78        self.dist = Distribution::NONE;
 79    }
 80}
 81
 82impl<D: DataMut, BRT: BlindRotationAlgo> FillUniform for BlindRotationKey<D, BRT> {
 83    fn fill_uniform(&mut self, source: &mut Source) {
 84        self.keys
 85            .iter_mut()
 86            .for_each(|key| key.fill_uniform(source));
 87    }
 88}
 89
 90impl<D: DataMut, BRT: BlindRotationAlgo> ReaderFrom for BlindRotationKey<D, BRT> {
 91    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
 92        match Distribution::read_from(reader) {
 93            Ok(dist) => self.dist = dist,
 94            Err(e) => return Err(e),
 95        }
 96        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
 97        if self.keys.len() != len {
 98            return Err(std::io::Error::new(
 99                std::io::ErrorKind::InvalidData,
100                format!("self.keys.len()={} != read len={}", self.keys.len(), len),
101            ));
102        }
103        for key in &mut self.keys {
104            key.read_from(reader)?;
105        }
106        Ok(())
107    }
108}
109
110impl<D: DataRef, BRT: BlindRotationAlgo> WriterTo for BlindRotationKey<D, BRT> {
111    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
112        match self.dist.write_to(writer) {
113            Ok(()) => {}
114            Err(e) => return Err(e),
115        }
116        writer.write_u64::<LittleEndian>(self.keys.len() as u64)?;
117        for key in &self.keys {
118            key.write_to(writer)?;
119        }
120        Ok(())
121    }
122}
123
124impl<D: DataRef, BRT: BlindRotationAlgo> BlindRotationKey<D, BRT> {
125    #[allow(dead_code)]
126    pub(crate) fn n(&self) -> usize {
127        self.keys[0].n()
128    }
129
130    #[allow(dead_code)]
131    pub(crate) fn rows(&self) -> usize {
132        self.keys[0].rows()
133    }
134
135    #[allow(dead_code)]
136    pub(crate) fn k(&self) -> usize {
137        self.keys[0].k()
138    }
139
140    #[allow(dead_code)]
141    pub(crate) fn size(&self) -> usize {
142        self.keys[0].size()
143    }
144
145    #[allow(dead_code)]
146    pub(crate) fn rank(&self) -> usize {
147        self.keys[0].rank()
148    }
149
150    pub(crate) fn basek(&self) -> usize {
151        self.keys[0].basek()
152    }
153
154    #[allow(dead_code)]
155    pub(crate) fn block_size(&self) -> usize {
156        match self.dist {
157            Distribution::BinaryBlock(value) => value,
158            _ => 1,
159        }
160    }
161}