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.iter_mut().for_each(|key| key.fill_uniform(log_bound, source));
 75    }
 76}
 77
 78impl<D: DataMut, BRT: BlindRotationAlgo> ReaderFrom for BlindRotationKeyCompressed<D, BRT> {
 79    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
 80        match Distribution::read_from(reader) {
 81            Ok(dist) => self.dist = dist,
 82            Err(e) => return Err(e),
 83        }
 84        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
 85        if self.keys.len() != len {
 86            return Err(std::io::Error::new(
 87                std::io::ErrorKind::InvalidData,
 88                format!("self.keys.len()={} != read len={}", self.keys.len(), len),
 89            ));
 90        }
 91        for key in &mut self.keys {
 92            key.read_from(reader)?;
 93        }
 94        Ok(())
 95    }
 96}
 97
 98impl<D: DataRef, BRT: BlindRotationAlgo> WriterTo for BlindRotationKeyCompressed<D, BRT> {
 99    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
100        match self.dist.write_to(writer) {
101            Ok(()) => {}
102            Err(e) => return Err(e),
103        }
104        writer.write_u64::<LittleEndian>(self.keys.len() as u64)?;
105        for key in &self.keys {
106            key.write_to(writer)?;
107        }
108        Ok(())
109    }
110}
111
112impl<D: DataRef, BRA: BlindRotationAlgo> BlindRotationKeyInfos for BlindRotationKeyCompressed<D, BRA> {
113    fn n_glwe(&self) -> Degree {
114        self.n()
115    }
116
117    fn n_lwe(&self) -> Degree {
118        Degree(self.keys.len() as u32)
119    }
120}
121
122impl<D: DataRef, BRA: BlindRotationAlgo> LWEInfos for BlindRotationKeyCompressed<D, BRA> {
123    fn n(&self) -> Degree {
124        self.keys[0].n()
125    }
126
127    fn size(&self) -> usize {
128        self.keys[0].size()
129    }
130
131    fn k(&self) -> TorusPrecision {
132        self.keys[0].k()
133    }
134
135    fn base2k(&self) -> Base2K {
136        self.keys[0].base2k()
137    }
138}
139
140impl<D: DataRef, BRA: BlindRotationAlgo> GLWEInfos for BlindRotationKeyCompressed<D, BRA> {
141    fn rank(&self) -> poulpy_core::layouts::Rank {
142        self.keys[0].rank()
143    }
144}
145
146impl<D: DataRef, BRA: BlindRotationAlgo> GGSWInfos for BlindRotationKeyCompressed<D, BRA> {
147    fn dnum(&self) -> poulpy_core::layouts::Dnum {
148        self.keys[0].dnum()
149    }
150
151    fn dsize(&self) -> poulpy_core::layouts::Dsize {
152        Dsize(1)
153    }
154}
155
156impl<D: DataRef, BRA: BlindRotationAlgo> BlindRotationKeyCompressed<D, BRA> {
157    #[allow(dead_code)]
158    pub(crate) fn block_size(&self) -> usize {
159        match self.dist {
160            Distribution::BinaryBlock(value) => value,
161            _ => 1,
162        }
163    }
164}