1use backend::hal::{
  2    api::{FillUniform, Reset, ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxAddScalarInplace, ZnxView, ZnxViewMut},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, ReaderFrom, ScalarZnx, ScalarZnxToRef, Scratch, WriterTo},
  4};
  5use sampling::source::Source;
  6
  7use crate::{Distribution, GGSWCiphertextCompressed, GGSWEncryptSkFamily, GLWESecretExec, Infos, LWESecret};
  8use std::fmt;
  9
 10#[derive(Clone)]
 11pub struct BlindRotationKeyCGGICompressed<D: Data> {
 12    pub(crate) keys: Vec<GGSWCiphertextCompressed<D>>,
 13    pub(crate) dist: Distribution,
 14}
 15
 16impl<D: DataRef> fmt::Debug for BlindRotationKeyCGGICompressed<D> {
 17    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 18        write!(f, "{}", self)
 19    }
 20}
 21
 22impl<D: Data> PartialEq for BlindRotationKeyCGGICompressed<D> {
 23    fn eq(&self, other: &Self) -> bool {
 24        if self.keys.len() != other.keys.len() {
 25            return false;
 26        }
 27        for (a, b) in self.keys.iter().zip(other.keys.iter()) {
 28            if a != b {
 29                return false;
 30            }
 31        }
 32        self.dist == other.dist
 33    }
 34}
 35
 36impl<D: Data> Eq for BlindRotationKeyCGGICompressed<D> {}
 37
 38impl<D: DataRef> fmt::Display for BlindRotationKeyCGGICompressed<D> {
 39    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 40        for (i, key) in self.keys.iter().enumerate() {
 41            write!(f, "key[{}]: {}", i, key)?;
 42        }
 43        writeln!(f, "{:?}", self.dist)
 44    }
 45}
 46
 47impl<D: DataMut> Reset for BlindRotationKeyCGGICompressed<D> {
 48    fn reset(&mut self) {
 49        self.keys.iter_mut().for_each(|key| key.reset());
 50        self.dist = Distribution::NONE;
 51    }
 52}
 53
 54impl<D: DataMut> FillUniform for BlindRotationKeyCGGICompressed<D> {
 55    fn fill_uniform(&mut self, source: &mut sampling::source::Source) {
 56        self.keys
 57            .iter_mut()
 58            .for_each(|key| key.fill_uniform(source));
 59    }
 60}
 61
 62use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 63
 64impl<D: DataMut> ReaderFrom for BlindRotationKeyCGGICompressed<D> {
 65    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
 66        match Distribution::read_from(reader) {
 67            Ok(dist) => self.dist = dist,
 68            Err(e) => return Err(e),
 69        }
 70        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
 71        if self.keys.len() != len {
 72            return Err(std::io::Error::new(
 73                std::io::ErrorKind::InvalidData,
 74                format!("self.keys.len()={} != read len={}", self.keys.len(), len),
 75            ));
 76        }
 77        for key in &mut self.keys {
 78            key.read_from(reader)?;
 79        }
 80        Ok(())
 81    }
 82}
 83
 84impl<D: DataRef> WriterTo for BlindRotationKeyCGGICompressed<D> {
 85    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
 86        match self.dist.write_to(writer) {
 87            Ok(()) => {}
 88            Err(e) => return Err(e),
 89        }
 90        writer.write_u64::<LittleEndian>(self.keys.len() as u64)?;
 91        for key in &self.keys {
 92            key.write_to(writer)?;
 93        }
 94        Ok(())
 95    }
 96}
 97
 98impl BlindRotationKeyCGGICompressed<Vec<u8>> {
 99    pub fn alloc(n_gglwe: usize, n_lwe: usize, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
100        let mut data: Vec<GGSWCiphertextCompressed<Vec<u8>>> = Vec::with_capacity(n_lwe);
101        (0..n_lwe).for_each(|_| {
102            data.push(GGSWCiphertextCompressed::alloc(
103                n_gglwe, basek, k, rows, 1, rank,
104            ))
105        });
106        Self {
107            keys: data,
108            dist: Distribution::NONE,
109        }
110    }
111
112    pub fn generate_from_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank: usize) -> usize
113    where
114        Module<B>: GGSWEncryptSkFamily<B>,
115    {
116        GGSWCiphertextCompressed::encrypt_sk_scratch_space(module, n, basek, k, rank)
117    }
118}
119
120impl<D: DataRef> BlindRotationKeyCGGICompressed<D> {
121    #[allow(dead_code)]
122    pub(crate) fn n(&self) -> usize {
123        self.keys[0].n()
124    }
125
126    #[allow(dead_code)]
127    pub(crate) fn rows(&self) -> usize {
128        self.keys[0].rows()
129    }
130
131    #[allow(dead_code)]
132    pub(crate) fn k(&self) -> usize {
133        self.keys[0].k()
134    }
135
136    #[allow(dead_code)]
137    pub(crate) fn size(&self) -> usize {
138        self.keys[0].size()
139    }
140
141    #[allow(dead_code)]
142    pub(crate) fn rank(&self) -> usize {
143        self.keys[0].rank()
144    }
145
146    #[allow(dead_code)]
147    pub(crate) fn basek(&self) -> usize {
148        self.keys[0].basek()
149    }
150
151    #[allow(dead_code)]
152    pub(crate) fn block_size(&self) -> usize {
153        match self.dist {
154            Distribution::BinaryBlock(value) => value,
155            _ => 1,
156        }
157    }
158}
159
160impl<D: DataMut> BlindRotationKeyCGGICompressed<D> {
161    pub fn generate_from_sk<DataSkGLWE, DataSkLWE, B: Backend>(
162        &mut self,
163        module: &Module<B>,
164        sk_glwe: &GLWESecretExec<DataSkGLWE, B>,
165        sk_lwe: &LWESecret<DataSkLWE>,
166        seed_xa: [u8; 32],
167        source_xe: &mut Source,
168        sigma: f64,
169        scratch: &mut Scratch<B>,
170    ) where
171        DataSkGLWE: DataRef,
172        DataSkLWE: DataRef,
173        Module<B>: GGSWEncryptSkFamily<B> + VecZnxAddScalarInplace,
174        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
175    {
176        #[cfg(debug_assertions)]
177        {
178            assert_eq!(self.keys.len(), sk_lwe.n());
179            assert!(sk_glwe.n() <= module.n());
180            assert_eq!(sk_glwe.rank(), self.keys[0].rank());
181            match sk_lwe.dist {
182                Distribution::BinaryBlock(_)
183                | Distribution::BinaryFixed(_)
184                | Distribution::BinaryProb(_)
185                | Distribution::ZERO => {}
186                _ => panic!(
187                    "invalid GLWESecret distribution: must be BinaryBlock, BinaryFixed or BinaryProb (or ZERO for debugging)"
188                ),
189            }
190        }
191
192        self.dist = sk_lwe.dist;
193
194        let mut pt: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(sk_glwe.n(), 1);
195        let sk_ref: ScalarZnx<&[u8]> = sk_lwe.data.to_ref();
196
197        let mut source_xa: Source = Source::new(seed_xa);
198
199        self.keys.iter_mut().enumerate().for_each(|(i, ggsw)| {
200            pt.at_mut(0, 0)[0] = sk_ref.at(0, 0)[i];
201            ggsw.encrypt_sk(
202                module,
203                &pt,
204                sk_glwe,
205                source_xa.new_seed(),
206                source_xe,
207                sigma,
208                scratch,
209            );
210        });
211    }
212}