1use backend::hal::{
  2    api::{
  3        MatZnxAlloc, ScalarZnxAlloc, ScratchAvailable, SvpPPolAlloc, SvpPrepare, TakeVecZnx, TakeVecZnxDft,
  4        VecZnxAddScalarInplace, VecZnxAllocBytes, ZnxView, ZnxViewMut,
  5    },
  6    layouts::{Backend, Data, DataMut, DataRef, Module, ReaderFrom, ScalarZnx, ScalarZnxToRef, Scratch, SvpPPol, WriterTo},
  7};
  8use sampling::source::Source;
  9
 10use crate::{
 11    Distribution, GGSWCiphertext, GGSWCiphertextExec, GGSWEncryptSkFamily, GGSWLayoutFamily, GLWESecretExec, Infos, LWESecret,
 12};
 13
 14pub struct BlindRotationKeyCGGI<D: Data> {
 15    pub(crate) keys: Vec<GGSWCiphertext<D>>,
 16    pub(crate) dist: Distribution,
 17}
 18
 19impl<D: Data> PartialEq for BlindRotationKeyCGGI<D> {
 20    fn eq(&self, other: &Self) -> bool {
 21        if self.keys.len() != other.keys.len() {
 22            return false;
 23        }
 24        for (a, b) in self.keys.iter().zip(other.keys.iter()) {
 25            if a != b {
 26                return false;
 27            }
 28        }
 29        self.dist == other.dist
 30    }
 31}
 32
 33impl<D: Data> Eq for BlindRotationKeyCGGI<D> {}
 34
 35use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 36
 37impl<D: DataMut> ReaderFrom for BlindRotationKeyCGGI<D> {
 38    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
 39        match Distribution::read_from(reader) {
 40            Ok(dist) => self.dist = dist,
 41            Err(e) => return Err(e),
 42        }
 43        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
 44        if self.keys.len() != len {
 45            return Err(std::io::Error::new(
 46                std::io::ErrorKind::InvalidData,
 47                format!("self.keys.len()={} != read len={}", self.keys.len(), len),
 48            ));
 49        }
 50        for key in &mut self.keys {
 51            key.read_from(reader)?;
 52        }
 53        Ok(())
 54    }
 55}
 56
 57impl<D: DataRef> WriterTo for BlindRotationKeyCGGI<D> {
 58    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
 59        match self.dist.write_to(writer) {
 60            Ok(()) => {}
 61            Err(e) => return Err(e),
 62        }
 63        writer.write_u64::<LittleEndian>(self.keys.len() as u64)?;
 64        for key in &self.keys {
 65            key.write_to(writer)?;
 66        }
 67        Ok(())
 68    }
 69}
 70
 71impl BlindRotationKeyCGGI<Vec<u8>> {
 72    pub fn alloc<B: Backend>(module: &Module<B>, n_lwe: usize, basek: usize, k: usize, rows: usize, rank: usize) -> Self
 73    where
 74        Module<B>: MatZnxAlloc,
 75    {
 76        let mut data: Vec<GGSWCiphertext<Vec<u8>>> = Vec::with_capacity(n_lwe);
 77        (0..n_lwe).for_each(|_| data.push(GGSWCiphertext::alloc(module, basek, k, rows, 1, rank)));
 78        Self {
 79            keys: data,
 80            dist: Distribution::NONE,
 81        }
 82    }
 83
 84    pub fn generate_from_sk_scratch_space<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize
 85    where
 86        Module<B>: GGSWEncryptSkFamily<B> + VecZnxAllocBytes,
 87    {
 88        GGSWCiphertext::encrypt_sk_scratch_space(module, basek, k, rank)
 89    }
 90}
 91
 92impl<D: DataRef> BlindRotationKeyCGGI<D> {
 93    #[allow(dead_code)]
 94    pub(crate) fn n(&self) -> usize {
 95        self.keys[0].n()
 96    }
 97
 98    #[allow(dead_code)]
 99    pub(crate) fn rows(&self) -> usize {
100        self.keys[0].rows()
101    }
102
103    #[allow(dead_code)]
104    pub(crate) fn k(&self) -> usize {
105        self.keys[0].k()
106    }
107
108    #[allow(dead_code)]
109    pub(crate) fn size(&self) -> usize {
110        self.keys[0].size()
111    }
112
113    #[allow(dead_code)]
114    pub(crate) fn rank(&self) -> usize {
115        self.keys[0].rank()
116    }
117
118    pub(crate) fn basek(&self) -> usize {
119        self.keys[0].basek()
120    }
121
122    #[allow(dead_code)]
123    pub(crate) fn block_size(&self) -> usize {
124        match self.dist {
125            Distribution::BinaryBlock(value) => value,
126            _ => 1,
127        }
128    }
129}
130
131impl<D: DataMut> BlindRotationKeyCGGI<D> {
132    pub fn generate_from_sk<DataSkGLWE, DataSkLWE, B: Backend>(
133        &mut self,
134        module: &Module<B>,
135        sk_glwe: &GLWESecretExec<DataSkGLWE, B>,
136        sk_lwe: &LWESecret<DataSkLWE>,
137        source_xa: &mut Source,
138        source_xe: &mut Source,
139        sigma: f64,
140        scratch: &mut Scratch<B>,
141    ) where
142        DataSkGLWE: DataRef,
143        DataSkLWE: DataRef,
144        Module<B>: GGSWEncryptSkFamily<B> + ScalarZnxAlloc + VecZnxAddScalarInplace,
145        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx<B>,
146    {
147        #[cfg(debug_assertions)]
148        {
149            assert_eq!(self.keys.len(), sk_lwe.n());
150            assert_eq!(sk_glwe.n(), module.n());
151            assert_eq!(sk_glwe.rank(), self.keys[0].rank());
152            match sk_lwe.dist {
153                Distribution::BinaryBlock(_)
154                | Distribution::BinaryFixed(_)
155                | Distribution::BinaryProb(_)
156                | Distribution::ZERO => {}
157                _ => panic!(
158                    "invalid GLWESecret distribution: must be BinaryBlock, BinaryFixed or BinaryProb (or ZERO for debugging)"
159                ),
160            }
161        }
162
163        self.dist = sk_lwe.dist;
164
165        let mut pt: ScalarZnx<Vec<u8>> = module.scalar_znx_alloc(1);
166        let sk_ref: ScalarZnx<&[u8]> = sk_lwe.data.to_ref();
167
168        self.keys.iter_mut().enumerate().for_each(|(i, ggsw)| {
169            pt.at_mut(0, 0)[0] = sk_ref.at(0, 0)[i];
170            ggsw.encrypt_sk(module, &pt, sk_glwe, source_xa, source_xe, sigma, scratch);
171        });
172    }
173}
174
175#[derive(PartialEq, Eq)]
176pub struct BlindRotationKeyCGGIExec<D: Data, B: Backend> {
177    pub(crate) data: Vec<GGSWCiphertextExec<D, B>>,
178    pub(crate) dist: Distribution,
179    pub(crate) x_pow_a: Option<Vec<SvpPPol<Vec<u8>, B>>>,
180}
181
182impl<D: Data, B: Backend> BlindRotationKeyCGGIExec<D, B> {
183    #[allow(dead_code)]
184    pub(crate) fn n(&self) -> usize {
185        self.data[0].n()
186    }
187
188    #[allow(dead_code)]
189    pub(crate) fn rows(&self) -> usize {
190        self.data[0].rows()
191    }
192
193    #[allow(dead_code)]
194    pub(crate) fn k(&self) -> usize {
195        self.data[0].k()
196    }
197
198    #[allow(dead_code)]
199    pub(crate) fn size(&self) -> usize {
200        self.data[0].size()
201    }
202
203    #[allow(dead_code)]
204    pub(crate) fn rank(&self) -> usize {
205        self.data[0].rank()
206    }
207
208    pub(crate) fn basek(&self) -> usize {
209        self.data[0].basek()
210    }
211
212    pub(crate) fn block_size(&self) -> usize {
213        match self.dist {
214            Distribution::BinaryBlock(value) => value,
215            _ => 1,
216        }
217    }
218}
219
220pub trait BlindRotationKeyCGGIExecLayoutFamily<B: Backend> = GGSWLayoutFamily<B> + SvpPPolAlloc<B> + SvpPrepare<B>;
221
222impl<B: Backend> BlindRotationKeyCGGIExec<Vec<u8>, B> {
223    pub fn alloc(module: &Module<B>, n_lwe: usize, basek: usize, k: usize, rows: usize, rank: usize) -> Self
224    where
225        Module<B>: BlindRotationKeyCGGIExecLayoutFamily<B>,
226    {
227        let mut data: Vec<GGSWCiphertextExec<Vec<u8>, B>> = Vec::with_capacity(n_lwe);
228        (0..n_lwe).for_each(|_| data.push(GGSWCiphertextExec::alloc(module, basek, k, rows, 1, rank)));
229        Self {
230            data,
231            dist: Distribution::NONE,
232            x_pow_a: None,
233        }
234    }
235
236    pub fn from<DataOther>(module: &Module<B>, other: &BlindRotationKeyCGGI<DataOther>, scratch: &mut Scratch<B>) -> Self
237    where
238        DataOther: DataRef,
239        Module<B>: BlindRotationKeyCGGIExecLayoutFamily<B> + ScalarZnxAlloc,
240    {
241        let mut brk: BlindRotationKeyCGGIExec<Vec<u8>, B> = Self::alloc(
242            module,
243            other.keys.len(),
244            other.basek(),
245            other.k(),
246            other.rows(),
247            other.rank(),
248        );
249        brk.prepare(module, other, scratch);
250        brk
251    }
252}
253
254impl<D: DataMut, B: Backend> BlindRotationKeyCGGIExec<D, B> {
255    pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &BlindRotationKeyCGGI<DataOther>, scratch: &mut Scratch<B>)
256    where
257        DataOther: DataRef,
258        Module<B>: BlindRotationKeyCGGIExecLayoutFamily<B> + ScalarZnxAlloc,
259    {
260        #[cfg(debug_assertions)]
261        {
262            assert_eq!(self.data.len(), other.keys.len());
263        }
264
265        self.data
266            .iter_mut()
267            .zip(other.keys.iter())
268            .for_each(|(ggsw_exec, other)| {
269                ggsw_exec.prepare(module, other, scratch);
270            });
271
272        self.dist = other.dist;
273
274        match other.dist {
275            Distribution::BinaryBlock(_) => {
276                let mut x_pow_a: Vec<SvpPPol<Vec<u8>, B>> = Vec::with_capacity(module.n() << 1);
277                let mut buf: ScalarZnx<Vec<u8>> = module.scalar_znx_alloc(1);
278                (0..module.n() << 1).for_each(|i| {
279                    let mut res: SvpPPol<Vec<u8>, B> = module.svp_ppol_alloc(1);
280                    set_xai_plus_y(module, i, 0, &mut res, &mut buf);
281                    x_pow_a.push(res);
282                });
283                self.x_pow_a = Some(x_pow_a);
284            }
285            _ => {}
286        }
287    }
288}
289
290pub fn set_xai_plus_y<A, C, B: Backend>(module: &Module<B>, ai: usize, y: i64, res: &mut SvpPPol<A, B>, buf: &mut ScalarZnx<C>)
291where
292    A: DataMut,
293    C: DataMut,
294    Module<B>: SvpPrepare<B>,
295{
296    let n: usize = module.n();
297
298    {
299        let raw: &mut [i64] = buf.at_mut(0, 0);
300        if ai < n {
301            raw[ai] = 1;
302        } else {
303            raw[(ai - n) & (n - 1)] = -1;
304        }
305        raw[0] += y;
306    }
307
308    module.svp_prepare(res, 0, buf, 0);
309
310    {
311        let raw: &mut [i64] = buf.at_mut(0, 0);
312
313        if ai < n {
314            raw[ai] = 0;
315        } else {
316            raw[(ai - n) & (n - 1)] = 0;
317        }
318        raw[0] = 0;
319    }
320}