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, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume,
7        VecZnxRshInplace, VmpApply, VmpApplyAdd, VmpApplyTmpBytes,
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        n: usize,
34        basek: usize,
35        out_k: usize,
36        in_k: usize,
37        ksk_k: usize,
38        digits: usize,
39        rank: usize,
40    ) -> usize
41    where
42        Module<B>: VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxBigNormalizeTmpBytes,
43    {
44        Self::automorphism_inplace_scratch_space(module, n, basek, out_k.min(in_k), ksk_k, digits, rank)
45    }
46
47    pub fn trace_inplace_scratch_space<B: Backend>(
48        module: &Module<B>,
49        n: usize,
50        basek: usize,
51        out_k: usize,
52        ksk_k: usize,
53        digits: usize,
54        rank: usize,
55    ) -> usize
56    where
57        Module<B>: VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxBigNormalizeTmpBytes,
58    {
59        Self::automorphism_inplace_scratch_space(module, n, basek, out_k, ksk_k, digits, rank)
60    }
61}
62
63impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
64    pub fn trace<DataLhs: DataRef, DataAK: DataRef, B: Backend>(
65        &mut self,
66        module: &Module<B>,
67        start: usize,
68        end: usize,
69        lhs: &GLWECiphertext<DataLhs>,
70        auto_keys: &HashMap<i64, GGLWEAutomorphismKeyPrepared<DataAK, B>>,
71        scratch: &mut Scratch<B>,
72    ) where
73        Module<B>: VecZnxDftAllocBytes
74            + VmpApplyTmpBytes
75            + VecZnxBigNormalizeTmpBytes
76            + VmpApply<B>
77            + VmpApplyAdd<B>
78            + VecZnxDftFromVecZnx<B>
79            + VecZnxDftToVecZnxBigConsume<B>
80            + VecZnxBigAddSmallInplace<B>
81            + VecZnxBigNormalize<B>
82            + VecZnxBigAutomorphismInplace<B>
83            + VecZnxRshInplace
84            + VecZnxCopy,
85        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
86    {
87        self.copy(module, lhs);
88        self.trace_inplace(module, start, end, auto_keys, scratch);
89    }
90
91    pub fn trace_inplace<DataAK: DataRef, B: Backend>(
92        &mut self,
93        module: &Module<B>,
94        start: usize,
95        end: usize,
96        auto_keys: &HashMap<i64, GGLWEAutomorphismKeyPrepared<DataAK, B>>,
97        scratch: &mut Scratch<B>,
98    ) where
99        Module<B>: VecZnxDftAllocBytes
100            + VmpApplyTmpBytes
101            + VecZnxBigNormalizeTmpBytes
102            + VmpApply<B>
103            + VmpApplyAdd<B>
104            + VecZnxDftFromVecZnx<B>
105            + VecZnxDftToVecZnxBigConsume<B>
106            + VecZnxBigAddSmallInplace<B>
107            + VecZnxBigNormalize<B>
108            + VecZnxBigAutomorphismInplace<B>
109            + VecZnxRshInplace,
110        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
111    {
112        (start..end).for_each(|i| {
113            self.rsh(module, 1);
114
115            let p: i64 = if i == 0 {
116                -1
117            } else {
118                module.galois_element(1 << (i - 1))
119            };
120
121            if let Some(key) = auto_keys.get(&p) {
122                self.automorphism_add_inplace(module, key, scratch);
123            } else {
124                panic!("auto_keys[{}] is empty", p)
125            }
126        });
127    }
128}
poulpy_core/glwe_trace.rs

poulpy_core/
glwe_trace.rs
