poulpy_core/
glwe_trace.rs

1use std::collections::HashMap;
2
3use poulpy_hal::{
4    api::{
5        ScratchAvailable, TakeVecZnxDft, VecZnxBigAddSmallInplace, VecZnxBigAutomorphismInplace, VecZnxBigNormalize,
6        VecZnxBigNormalizeTmpBytes, VecZnxCopy, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume, VecZnxRshInplace,
7        VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
8    },
9    layouts::{Backend, DataMut, DataRef, Module, Scratch},
10};
11
12use crate::{
13    layouts::{GLWECiphertext, prepared::GGLWEAutomorphismKeyPrepared},
14    operations::GLWEOperations,
15};
16
17impl GLWECiphertext<Vec<u8>> {
18    pub fn trace_galois_elements<B: Backend>(module: &Module<B>) -> Vec<i64> {
19        let mut gal_els: Vec<i64> = Vec::new();
20        (0..module.log_n()).for_each(|i| {
21            if i == 0 {
22                gal_els.push(-1);
23            } else {
24                gal_els.push(module.galois_element(1 << (i - 1)));
25            }
26        });
27        gal_els
28    }
29
30    #[allow(clippy::too_many_arguments)]
31    pub fn trace_scratch_space<B: Backend>(
32        module: &Module<B>,
33        basek: usize,
34        out_k: usize,
35        in_k: usize,
36        ksk_k: usize,
37        digits: usize,
38        rank: usize,
39    ) -> usize
40    where
41        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes,
42    {
43        Self::automorphism_inplace_scratch_space(module, basek, out_k.min(in_k), ksk_k, digits, rank)
44    }
45
46    pub fn trace_inplace_scratch_space<B: Backend>(
47        module: &Module<B>,
48        basek: usize,
49        out_k: usize,
50        ksk_k: usize,
51        digits: usize,
52        rank: usize,
53    ) -> usize
54    where
55        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes,
56    {
57        Self::automorphism_inplace_scratch_space(module, basek, out_k, ksk_k, digits, rank)
58    }
59}
60
61impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
62    pub fn trace<DataLhs: DataRef, DataAK: DataRef, B: Backend>(
63        &mut self,
64        module: &Module<B>,
65        start: usize,
66        end: usize,
67        lhs: &GLWECiphertext<DataLhs>,
68        auto_keys: &HashMap<i64, GGLWEAutomorphismKeyPrepared<DataAK, B>>,
69        scratch: &mut Scratch<B>,
70    ) where
71        Module<B>: VecZnxDftAllocBytes
72            + VmpApplyDftToDftTmpBytes
73            + VecZnxBigNormalizeTmpBytes
74            + VmpApplyDftToDft<B>
75            + VmpApplyDftToDftAdd<B>
76            + VecZnxDftApply<B>
77            + VecZnxIdftApplyConsume<B>
78            + VecZnxBigAddSmallInplace<B>
79            + VecZnxBigNormalize<B>
80            + VecZnxBigAutomorphismInplace<B>
81            + VecZnxRshInplace<B>
82            + VecZnxCopy,
83        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
84    {
85        self.copy(module, lhs);
86        self.trace_inplace(module, start, end, auto_keys, scratch);
87    }
88
89    pub fn trace_inplace<DataAK: DataRef, B: Backend>(
90        &mut self,
91        module: &Module<B>,
92        start: usize,
93        end: usize,
94        auto_keys: &HashMap<i64, GGLWEAutomorphismKeyPrepared<DataAK, B>>,
95        scratch: &mut Scratch<B>,
96    ) where
97        Module<B>: VecZnxDftAllocBytes
98            + VmpApplyDftToDftTmpBytes
99            + VecZnxBigNormalizeTmpBytes
100            + VmpApplyDftToDft<B>
101            + VmpApplyDftToDftAdd<B>
102            + VecZnxDftApply<B>
103            + VecZnxIdftApplyConsume<B>
104            + VecZnxBigAddSmallInplace<B>
105            + VecZnxBigNormalize<B>
106            + VecZnxBigAutomorphismInplace<B>
107            + VecZnxRshInplace<B>,
108        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
109    {
110        (start..end).for_each(|i| {
111            self.rsh(module, 1, scratch);
112
113            let p: i64 = if i == 0 {
114                -1
115            } else {
116                module.galois_element(1 << (i - 1))
117            };
118
119            if let Some(key) = auto_keys.get(&p) {
120                self.automorphism_add_inplace(module, key, scratch);
121            } else {
122                panic!("auto_keys[{}] is empty", p)
123            }
124        });
125    }
126}
poulpy_core/glwe_trace.rs

poulpy_core/
glwe_trace.rs
