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, Degree, Dsize, GGSWInfos, GLWEInfos, LWEInfos, TorusPrecision, compressed::GGSWCompressed},
 12};
 13
 14use crate::bin_fhe::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
 23pub trait BlindRotationKeyCompressedFactory<BRA: BlindRotationAlgo> {
 24    fn blind_rotation_key_compressed_alloc<A>(infos: &A) -> BlindRotationKeyCompressed<Vec<u8>, BRA>
 25    where
 26        A: BlindRotationKeyInfos;
 27}
 28
 29impl<BRA: BlindRotationAlgo> BlindRotationKeyCompressed<Vec<u8>, BRA>
 30where
 31    Self: BlindRotationKeyCompressedFactory<BRA>,
 32{
 33    pub fn alloc<A>(infos: &A) -> BlindRotationKeyCompressed<Vec<u8>, BRA>
 34    where
 35        A: BlindRotationKeyInfos,
 36    {
 37        Self::blind_rotation_key_compressed_alloc(infos)
 38    }
 39}
 40
 41impl<D: DataRef, BRT: BlindRotationAlgo> fmt::Debug for BlindRotationKeyCompressed<D, BRT> {
 42    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 43        write!(f, "{self}")
 44    }
 45}
 46
 47impl<D: Data, BRT: BlindRotationAlgo> PartialEq for BlindRotationKeyCompressed<D, BRT> {
 48    fn eq(&self, other: &Self) -> bool {
 49        if self.keys.len() != other.keys.len() {
 50            return false;
 51        }
 52        for (a, b) in self.keys.iter().zip(other.keys.iter()) {
 53            if a != b {
 54                return false;
 55            }
 56        }
 57        self.dist == other.dist && self._phantom == other._phantom
 58    }
 59}
 60
 61impl<D: Data, BRT: BlindRotationAlgo> Eq for BlindRotationKeyCompressed<D, BRT> {}
 62
 63impl<D: DataRef, BRT: BlindRotationAlgo> fmt::Display for BlindRotationKeyCompressed<D, BRT> {
 64    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 65        for (i, key) in self.keys.iter().enumerate() {
 66            write!(f, "key[{i}]: {key}")?;
 67        }
 68        writeln!(f, "{:?}", self.dist)
 69    }
 70}
 71
 72impl<D: DataMut, BRT: BlindRotationAlgo> FillUniform for BlindRotationKeyCompressed<D, BRT> {
 73    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
 74        self.keys
 75            .iter_mut()
 76            .for_each(|key| key.fill_uniform(log_bound, source));
 77    }
 78}
 79
 80impl<D: DataMut, BRT: BlindRotationAlgo> ReaderFrom for BlindRotationKeyCompressed<D, BRT> {
 81    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
 82        match Distribution::read_from(reader) {
 83            Ok(dist) => self.dist = dist,
 84            Err(e) => return Err(e),
 85        }
 86        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
 87        if self.keys.len() != len {
 88            return Err(std::io::Error::new(
 89                std::io::ErrorKind::InvalidData,
 90                format!("self.keys.len()={} != read len={}", self.keys.len(), len),
 91            ));
 92        }
 93        for key in &mut self.keys {
 94            key.read_from(reader)?;
 95        }
 96        Ok(())
 97    }
 98}
 99
100impl<D: DataRef, BRT: BlindRotationAlgo> WriterTo for BlindRotationKeyCompressed<D, BRT> {
101    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
102        match self.dist.write_to(writer) {
103            Ok(()) => {}
104            Err(e) => return Err(e),
105        }
106        writer.write_u64::<LittleEndian>(self.keys.len() as u64)?;
107        for key in &self.keys {
108            key.write_to(writer)?;
109        }
110        Ok(())
111    }
112}
113
114impl<D: DataRef, BRA: BlindRotationAlgo> BlindRotationKeyInfos for BlindRotationKeyCompressed<D, BRA> {
115    fn n_glwe(&self) -> Degree {
116        self.n()
117    }
118
119    fn n_lwe(&self) -> Degree {
120        Degree(self.keys.len() as u32)
121    }
122}
123
124impl<D: DataRef, BRA: BlindRotationAlgo> LWEInfos for BlindRotationKeyCompressed<D, BRA> {
125    fn n(&self) -> Degree {
126        self.keys[0].n()
127    }
128
129    fn size(&self) -> usize {
130        self.keys[0].size()
131    }
132
133    fn k(&self) -> TorusPrecision {
134        self.keys[0].k()
135    }
136
137    fn base2k(&self) -> Base2K {
138        self.keys[0].base2k()
139    }
140}
141
142impl<D: DataRef, BRA: BlindRotationAlgo> GLWEInfos for BlindRotationKeyCompressed<D, BRA> {
143    fn rank(&self) -> poulpy_core::layouts::Rank {
144        self.keys[0].rank()
145    }
146}
147
148impl<D: DataRef, BRA: BlindRotationAlgo> GGSWInfos for BlindRotationKeyCompressed<D, BRA> {
149    fn dnum(&self) -> poulpy_core::layouts::Dnum {
150        self.keys[0].dnum()
151    }
152
153    fn dsize(&self) -> poulpy_core::layouts::Dsize {
154        Dsize(1)
155    }
156}
157
158impl<D: DataRef, BRA: BlindRotationAlgo> BlindRotationKeyCompressed<D, BRA> {
159    #[allow(dead_code)]
160    pub(crate) fn block_size(&self) -> usize {
161        match self.dist {
162            Distribution::BinaryBlock(value) => value,
163            _ => 1,
164        }
165    }
166}