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