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