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