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