1use poulpy_hal::{
  2    layouts::{Data, DataMut, DataRef, FillUniform, ReaderFrom, WriterTo},
  3    source::Source,
  4};
  5
  6use std::{fmt, marker::PhantomData};
  7
  8use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
  9use poulpy_core::{
 10    Distribution,
 11    layouts::{Base2K, Dsize, GGSWInfos, GLWEInfos, LWEInfos, RingDegree, TorusPrecision, compressed::GGSWCompressed},
 12};
 13
 14use crate::tfhe::blind_rotation::{BlindRotationAlgo, BlindRotationKeyInfos};
 15
 16#[derive(Clone)]
 17pub struct BlindRotationKeyCompressed<D: Data, BRT: BlindRotationAlgo> {
 18    pub(crate) keys: Vec<GGSWCompressed<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> FillUniform for BlindRotationKeyCompressed<D, BRT> {
 55    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
 56        self.keys
 57            .iter_mut()
 58            .for_each(|key| key.fill_uniform(log_bound, source));
 59    }
 60}
 61
 62impl<D: DataMut, BRT: BlindRotationAlgo> ReaderFrom for BlindRotationKeyCompressed<D, BRT> {
 63    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
 64        match Distribution::read_from(reader) {
 65            Ok(dist) => self.dist = dist,
 66            Err(e) => return Err(e),
 67        }
 68        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
 69        if self.keys.len() != len {
 70            return Err(std::io::Error::new(
 71                std::io::ErrorKind::InvalidData,
 72                format!("self.keys.len()={} != read len={}", self.keys.len(), len),
 73            ));
 74        }
 75        for key in &mut self.keys {
 76            key.read_from(reader)?;
 77        }
 78        Ok(())
 79    }
 80}
 81
 82impl<D: DataRef, BRT: BlindRotationAlgo> WriterTo for BlindRotationKeyCompressed<D, BRT> {
 83    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
 84        match self.dist.write_to(writer) {
 85            Ok(()) => {}
 86            Err(e) => return Err(e),
 87        }
 88        writer.write_u64::<LittleEndian>(self.keys.len() as u64)?;
 89        for key in &self.keys {
 90            key.write_to(writer)?;
 91        }
 92        Ok(())
 93    }
 94}
 95
 96impl<D: DataRef, BRA: BlindRotationAlgo> BlindRotationKeyInfos for BlindRotationKeyCompressed<D, BRA> {
 97    fn n_glwe(&self) -> RingDegree {
 98        self.n()
 99    }
100
101    fn n_lwe(&self) -> RingDegree {
102        RingDegree(self.keys.len() as u32)
103    }
104}
105
106impl<D: DataRef, BRA: BlindRotationAlgo> LWEInfos for BlindRotationKeyCompressed<D, BRA> {
107    fn n(&self) -> RingDegree {
108        self.keys[0].n()
109    }
110
111    fn size(&self) -> usize {
112        self.keys[0].size()
113    }
114
115    fn k(&self) -> TorusPrecision {
116        self.keys[0].k()
117    }
118
119    fn base2k(&self) -> Base2K {
120        self.keys[0].base2k()
121    }
122}
123
124impl<D: DataRef, BRA: BlindRotationAlgo> GLWEInfos for BlindRotationKeyCompressed<D, BRA> {
125    fn rank(&self) -> poulpy_core::layouts::Rank {
126        self.keys[0].rank()
127    }
128}
129
130impl<D: DataRef, BRA: BlindRotationAlgo> GGSWInfos for BlindRotationKeyCompressed<D, BRA> {
131    fn dnum(&self) -> poulpy_core::layouts::Dnum {
132        self.keys[0].dnum()
133    }
134
135    fn dsize(&self) -> poulpy_core::layouts::Dsize {
136        Dsize(1)
137    }
138}
139
140impl<D: DataRef, BRA: BlindRotationAlgo> BlindRotationKeyCompressed<D, BRA> {
141    #[allow(dead_code)]
142    pub(crate) fn block_size(&self) -> usize {
143        match self.dist {
144            Distribution::BinaryBlock(value) => value,
145            _ => 1,
146        }
147    }
148}