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