1use backend::hal::{
  2    api::{
  3        ScratchAvailable, SvpPPolAlloc, SvpPrepare, TakeVecZnx, TakeVecZnxDft, VecZnxAddScalarInplace, ZnxInfos, ZnxView,
  4        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(n_gglwe: usize, n_lwe: usize, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
 73        let mut data: Vec<GGSWCiphertext<Vec<u8>>> = Vec::with_capacity(n_lwe);
 74        (0..n_lwe).for_each(|_| data.push(GGSWCiphertext::alloc(n_gglwe, basek, k, rows, 1, rank)));
 75        Self {
 76            keys: data,
 77            dist: Distribution::NONE,
 78        }
 79    }
 80
 81    pub fn generate_from_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank: usize) -> usize
 82    where
 83        Module<B>: GGSWEncryptSkFamily<B>,
 84    {
 85        GGSWCiphertext::encrypt_sk_scratch_space(module, n, basek, k, rank)
 86    }
 87}
 88
 89impl<D: DataRef> BlindRotationKeyCGGI<D> {
 90    #[allow(dead_code)]
 91    pub(crate) fn n(&self) -> usize {
 92        self.keys[0].n()
 93    }
 94
 95    #[allow(dead_code)]
 96    pub(crate) fn rows(&self) -> usize {
 97        self.keys[0].rows()
 98    }
 99
100    #[allow(dead_code)]
101    pub(crate) fn k(&self) -> usize {
102        self.keys[0].k()
103    }
104
105    #[allow(dead_code)]
106    pub(crate) fn size(&self) -> usize {
107        self.keys[0].size()
108    }
109
110    #[allow(dead_code)]
111    pub(crate) fn rank(&self) -> usize {
112        self.keys[0].rank()
113    }
114
115    pub(crate) fn basek(&self) -> usize {
116        self.keys[0].basek()
117    }
118
119    #[allow(dead_code)]
120    pub(crate) fn block_size(&self) -> usize {
121        match self.dist {
122            Distribution::BinaryBlock(value) => value,
123            _ => 1,
124        }
125    }
126}
127
128impl<D: DataMut> BlindRotationKeyCGGI<D> {
129    pub fn generate_from_sk<DataSkGLWE, DataSkLWE, B: Backend>(
130        &mut self,
131        module: &Module<B>,
132        sk_glwe: &GLWESecretExec<DataSkGLWE, B>,
133        sk_lwe: &LWESecret<DataSkLWE>,
134        source_xa: &mut Source,
135        source_xe: &mut Source,
136        sigma: f64,
137        scratch: &mut Scratch<B>,
138    ) where
139        DataSkGLWE: DataRef,
140        DataSkLWE: DataRef,
141        Module<B>: GGSWEncryptSkFamily<B> + VecZnxAddScalarInplace,
142        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
143    {
144        #[cfg(debug_assertions)]
145        {
146            assert_eq!(self.keys.len(), sk_lwe.n());
147            assert!(sk_glwe.n() <= module.n());
148            assert_eq!(sk_glwe.rank(), self.keys[0].rank());
149            match sk_lwe.dist {
150                Distribution::BinaryBlock(_)
151                | Distribution::BinaryFixed(_)
152                | Distribution::BinaryProb(_)
153                | Distribution::ZERO => {}
154                _ => panic!(
155                    "invalid GLWESecret distribution: must be BinaryBlock, BinaryFixed or BinaryProb (or ZERO for debugging)"
156                ),
157            }
158        }
159
160        self.dist = sk_lwe.dist;
161
162        let mut pt: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(sk_glwe.n(), 1);
163        let sk_ref: ScalarZnx<&[u8]> = sk_lwe.data.to_ref();
164
165        self.keys.iter_mut().enumerate().for_each(|(i, ggsw)| {
166            pt.at_mut(0, 0)[0] = sk_ref.at(0, 0)[i];
167            ggsw.encrypt_sk(module, &pt, sk_glwe, source_xa, source_xe, sigma, scratch);
168        });
169    }
170}
171
172#[derive(PartialEq, Eq)]
173pub struct BlindRotationKeyCGGIExec<D: Data, B: Backend> {
174    pub(crate) data: Vec<GGSWCiphertextExec<D, B>>,
175    pub(crate) dist: Distribution,
176    pub(crate) x_pow_a: Option<Vec<SvpPPol<Vec<u8>, B>>>,
177}
178
179impl<D: Data, B: Backend> BlindRotationKeyCGGIExec<D, B> {
180    #[allow(dead_code)]
181    pub(crate) fn n(&self) -> usize {
182        self.data[0].n()
183    }
184
185    #[allow(dead_code)]
186    pub(crate) fn rows(&self) -> usize {
187        self.data[0].rows()
188    }
189
190    #[allow(dead_code)]
191    pub(crate) fn k(&self) -> usize {
192        self.data[0].k()
193    }
194
195    #[allow(dead_code)]
196    pub(crate) fn size(&self) -> usize {
197        self.data[0].size()
198    }
199
200    #[allow(dead_code)]
201    pub(crate) fn rank(&self) -> usize {
202        self.data[0].rank()
203    }
204
205    pub(crate) fn basek(&self) -> usize {
206        self.data[0].basek()
207    }
208
209    pub(crate) fn block_size(&self) -> usize {
210        match self.dist {
211            Distribution::BinaryBlock(value) => value,
212            _ => 1,
213        }
214    }
215}
216
217pub trait BlindRotationKeyCGGIExecLayoutFamily<B: Backend> = GGSWLayoutFamily<B> + SvpPPolAlloc<B> + SvpPrepare<B>;
218
219impl<B: Backend> BlindRotationKeyCGGIExec<Vec<u8>, B> {
220    pub fn alloc(module: &Module<B>, n_glwe: usize, n_lwe: usize, basek: usize, k: usize, rows: usize, rank: usize) -> Self
221    where
222        Module<B>: BlindRotationKeyCGGIExecLayoutFamily<B>,
223    {
224        let mut data: Vec<GGSWCiphertextExec<Vec<u8>, B>> = Vec::with_capacity(n_lwe);
225        (0..n_lwe).for_each(|_| {
226            data.push(GGSWCiphertextExec::alloc(
227                module, n_glwe, basek, k, rows, 1, rank,
228            ))
229        });
230        Self {
231            data,
232            dist: Distribution::NONE,
233            x_pow_a: None,
234        }
235    }
236
237    pub fn from<DataOther>(module: &Module<B>, other: &BlindRotationKeyCGGI<DataOther>, scratch: &mut Scratch<B>) -> Self
238    where
239        DataOther: DataRef,
240        Module<B>: BlindRotationKeyCGGIExecLayoutFamily<B>,
241    {
242        let mut brk: BlindRotationKeyCGGIExec<Vec<u8>, B> = Self::alloc(
243            module,
244            other.n(),
245            other.keys.len(),
246            other.basek(),
247            other.k(),
248            other.rows(),
249            other.rank(),
250        );
251        brk.prepare(module, other, scratch);
252        brk
253    }
254}
255
256impl<D: DataMut, B: Backend> BlindRotationKeyCGGIExec<D, B> {
257    pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &BlindRotationKeyCGGI<DataOther>, scratch: &mut Scratch<B>)
258    where
259        DataOther: DataRef,
260        Module<B>: BlindRotationKeyCGGIExecLayoutFamily<B>,
261    {
262        #[cfg(debug_assertions)]
263        {
264            assert_eq!(self.data.len(), other.keys.len());
265        }
266
267        let n: usize = other.n();
268
269        self.data
270            .iter_mut()
271            .zip(other.keys.iter())
272            .for_each(|(ggsw_exec, other)| {
273                ggsw_exec.prepare(module, other, scratch);
274            });
275
276        self.dist = other.dist;
277
278        match other.dist {
279            Distribution::BinaryBlock(_) => {
280                let mut x_pow_a: Vec<SvpPPol<Vec<u8>, B>> = Vec::with_capacity(n << 1);
281                let mut buf: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
282                (0..n << 1).for_each(|i| {
283                    let mut res: SvpPPol<Vec<u8>, B> = module.svp_ppol_alloc(n, 1);
284                    set_xai_plus_y(module, i, 0, &mut res, &mut buf);
285                    x_pow_a.push(res);
286                });
287                self.x_pow_a = Some(x_pow_a);
288            }
289            _ => {}
290        }
291    }
292}
293
294pub 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>)
295where
296    A: DataMut,
297    C: DataMut,
298    Module<B>: SvpPrepare<B>,
299{
300    let n: usize = res.n();
301
302    {
303        let raw: &mut [i64] = buf.at_mut(0, 0);
304        if ai < n {
305            raw[ai] = 1;
306        } else {
307            raw[(ai - n) & (n - 1)] = -1;
308        }
309        raw[0] += y;
310    }
311
312    module.svp_prepare(res, 0, buf, 0);
313
314    {
315        let raw: &mut [i64] = buf.at_mut(0, 0);
316
317        if ai < n {
318            raw[ai] = 0;
319        } else {
320            raw[(ai - n) & (n - 1)] = 0;
321        }
322        raw[0] = 0;
323    }
324}