poulpy/core/src/glwe/packing.rs
core/src/glwe/packing.rsRaw
1use crate::{GLWEAutomorphismKey, GLWECiphertext, GLWEOps, Infos, ScratchCore};
2use std::collections::HashMap;
3
4use backend::{FFT64, Module, Scratch, VecZnxOps};
5
6/// [StreamPacker] enables only the fly GLWE packing
7/// with constant memory of Log(N) ciphertexts.
8/// Main difference with usual GLWE packing is that
9/// the output is bit-reversed.
10pub struct GLWEPacker {
11 accumulators: Vec<Accumulator>,
12 log_batch: usize,
13 counter: usize,
14}
15
16/// [Accumulator] stores intermediate packing result.
17/// There are Log(N) such accumulators in a [StreamPacker].
18struct Accumulator {
19 data: GLWECiphertext<Vec<u8>>,
20 value: bool, // Implicit flag for zero ciphertext
21 control: bool, // Can be combined with incoming value
22}
23
24impl Accumulator {
25 /// Allocates a new [Accumulator].
26 ///
27 /// #Arguments
28 ///
29 /// * `module`: static backend FFT tables.
30 /// * `basek`: base 2 logarithm of the GLWE ciphertext in memory digit representation.
31 /// * `k`: base 2 precision of the GLWE ciphertext precision over the Torus.
32 /// * `rank`: rank of the GLWE ciphertext.
33 pub fn alloc(module: &Module<FFT64>, basek: usize, k: usize, rank: usize) -> Self {
34 Self {
35 data: GLWECiphertext::alloc(module, basek, k, rank),
36 value: false,
37 control: false,
38 }
39 }
40}
41
42impl GLWEPacker {
43 /// Instantiates a new [StreamPacker].
44 ///
45 /// #Arguments
46 ///
47 /// * `module`: static backend FFT tables.
48 /// * `log_batch`: packs coefficients which are multiples of X^{N/2^log_batch}.
49 /// i.e. with `log_batch=0` only the constant coefficient is packed
50 /// and N GLWE ciphertext can be packed. With `log_batch=2` all coefficients
51 /// which are multiples of X^{N/4} are packed. Meaning that N/4 ciphertexts
52 /// can be packed.
53 /// * `basek`: base 2 logarithm of the GLWE ciphertext in memory digit representation.
54 /// * `k`: base 2 precision of the GLWE ciphertext precision over the Torus.
55 /// * `rank`: rank of the GLWE ciphertext.
56 pub fn new(module: &Module<FFT64>, log_batch: usize, basek: usize, k: usize, rank: usize) -> Self {
57 let mut accumulators: Vec<Accumulator> = Vec::<Accumulator>::new();
58 let log_n: usize = module.log_n();
59 (0..log_n - log_batch).for_each(|_| accumulators.push(Accumulator::alloc(module, basek, k, rank)));
60 Self {
61 accumulators: accumulators,
62 log_batch,
63 counter: 0,
64 }
65 }
66
67 /// Implicit reset of the internal state (to be called before a new packing procedure).
68 fn reset(&mut self) {
69 for i in 0..self.accumulators.len() {
70 self.accumulators[i].value = false;
71 self.accumulators[i].control = false;
72 }
73 self.counter = 0;
74 }
75
76 /// Number of scratch space bytes required to call [Self::add].
77 pub fn scratch_space(module: &Module<FFT64>, basek: usize, ct_k: usize, k_ksk: usize, digits: usize, rank: usize) -> usize {
78 pack_core_scratch_space(module, basek, ct_k, k_ksk, digits, rank)
79 }
80
81 pub fn galois_elements(module: &Module<FFT64>) -> Vec<i64> {
82 GLWECiphertext::trace_galois_elements(module)
83 }
84
85 /// Adds a GLWE ciphertext to the [StreamPacker].
86 /// #Arguments
87 ///
88 /// * `module`: static backend FFT tables.
89 /// * `res`: space to append fully packed ciphertext. Only when the number
90 /// of packed ciphertexts reaches N/2^log_batch is a result written.
91 /// * `a`: ciphertext to pack. Can optionally give None to pack a 0 ciphertext.
92 /// * `auto_keys`: a [HashMap] containing the [AutomorphismKey]s.
93 /// * `scratch`: scratch space of size at least [Self::add_scratch_space].
94 pub fn add<DataA: AsRef<[u8]>, DataAK: AsRef<[u8]>>(
95 &mut self,
96 module: &Module<FFT64>,
97 a: Option<&GLWECiphertext<DataA>>,
98 auto_keys: &HashMap<i64, GLWEAutomorphismKey<DataAK, FFT64>>,
99 scratch: &mut Scratch,
100 ) {
101 assert!(
102 self.counter < module.n(),
103 "Packing limit of {} reached",
104 module.n() >> self.log_batch
105 );
106
107 pack_core(
108 module,
109 a,
110 &mut self.accumulators,
111 self.log_batch,
112 auto_keys,
113 scratch,
114 );
115 self.counter += 1 << self.log_batch;
116 }
117
118 /// Flush result to`res`.
119 pub fn flush<Data: AsMut<[u8]> + AsRef<[u8]>>(&mut self, module: &Module<FFT64>, res: &mut GLWECiphertext<Data>) {
120 assert!(self.counter == module.n());
121 // Copy result GLWE into res GLWE
122 res.copy(
123 module,
124 &self.accumulators[module.log_n() - self.log_batch - 1].data,
125 );
126
127 self.reset();
128 }
129}
130
131fn pack_core_scratch_space(module: &Module<FFT64>, basek: usize, ct_k: usize, k_ksk: usize, digits: usize, rank: usize) -> usize {
132 combine_scratch_space(module, basek, ct_k, k_ksk, digits, rank)
133}
134
135fn pack_core<D: AsRef<[u8]>, DataAK: AsRef<[u8]>>(
136 module: &Module<FFT64>,
137 a: Option<&GLWECiphertext<D>>,
138 accumulators: &mut [Accumulator],
139 i: usize,
140 auto_keys: &HashMap<i64, GLWEAutomorphismKey<DataAK, FFT64>>,
141 scratch: &mut Scratch,
142) {
143 let log_n: usize = module.log_n();
144
145 if i == log_n {
146 return;
147 }
148
149 // Isolate the first accumulator
150 let (acc_prev, acc_next) = accumulators.split_at_mut(1);
151
152 // Control = true accumlator is free to overide
153 if !acc_prev[0].control {
154 let acc_mut_ref: &mut Accumulator = &mut acc_prev[0]; // from split_at_mut
155
156 // No previous value -> copies and sets flags accordingly
157 if let Some(a_ref) = a {
158 acc_mut_ref.data.copy(module, a_ref);
159 acc_mut_ref.value = true
160 } else {
161 acc_mut_ref.value = false
162 }
163 acc_mut_ref.control = true; // Able to be combined on next call
164 } else {
165 // Compresses acc_prev <- combine(acc_prev, a).
166 combine(module, &mut acc_prev[0], a, i, auto_keys, scratch);
167 acc_prev[0].control = false;
168
169 // Propagates to next accumulator
170 if acc_prev[0].value {
171 pack_core(
172 module,
173 Some(&acc_prev[0].data),
174 acc_next,
175 i + 1,
176 auto_keys,
177 scratch,
178 );
179 } else {
180 pack_core(
181 module,
182 None::<&GLWECiphertext<Vec<u8>>>,
183 acc_next,
184 i + 1,
185 auto_keys,
186 scratch,
187 );
188 }
189 }
190}
191
192fn combine_scratch_space(module: &Module<FFT64>, basek: usize, ct_k: usize, k_ksk: usize, digits: usize, rank: usize) -> usize {
193 GLWECiphertext::bytes_of(module, basek, ct_k, rank)
194 + (GLWECiphertext::rsh_scratch_space(module)
195 | GLWECiphertext::automorphism_scratch_space(module, basek, ct_k, ct_k, k_ksk, digits, rank))
196}
197
198/// [combine] merges two ciphertexts together.
199fn combine<D: AsRef<[u8]>, DataAK: AsRef<[u8]>>(
200 module: &Module<FFT64>,
201 acc: &mut Accumulator,
202 b: Option<&GLWECiphertext<D>>,
203 i: usize,
204 auto_keys: &HashMap<i64, GLWEAutomorphismKey<DataAK, FFT64>>,
205 scratch: &mut Scratch,
206) {
207 let log_n: usize = module.log_n();
208 let a: &mut GLWECiphertext<Vec<u8>> = &mut acc.data;
209 let basek: usize = a.basek();
210 let k: usize = a.k();
211 let rank: usize = a.rank();
212
213 let gal_el: i64;
214
215 if i == 0 {
216 gal_el = -1;
217 } else {
218 gal_el = module.galois_element(1 << (i - 1))
219 }
220
221 let t: i64 = 1 << (log_n - i - 1);
222
223 // Goal is to evaluate: a = a + b*X^t + phi(a - b*X^t))
224 // We also use the identity: AUTO(a * X^t, g) = -X^t * AUTO(a, g)
225 // where t = 2^(log_n - i - 1) and g = 5^{2^(i - 1)}
226 // Different cases for wether a and/or b are zero.
227 //
228 // Implicite RSH without modulus switch, introduces extra I(X) * Q/2 on decryption.
229 // Necessary so that the scaling of the plaintext remains constant.
230 // It however is ok to do so here because coefficients are eventually
231 // either mapped to garbage or twice their value which vanishes I(X)
232 // since 2*(I(X) * Q/2) = I(X) * Q = 0 mod Q.
233 if acc.value {
234 if let Some(b) = b {
235 let (mut tmp_b, scratch_1) = scratch.tmp_glwe_ct(module, basek, k, rank);
236
237 // a = a * X^-t
238 a.rotate_inplace(module, -t);
239
240 // tmp_b = a * X^-t - b
241 tmp_b.sub(module, a, b);
242 tmp_b.rsh(1, scratch_1);
243
244 // a = a * X^-t + b
245 a.add_inplace(module, b);
246 a.rsh(1, scratch_1);
247
248 tmp_b.normalize_inplace(module, scratch_1);
249
250 // tmp_b = phi(a * X^-t - b)
251 if let Some(key) = auto_keys.get(&gal_el) {
252 tmp_b.automorphism_inplace(module, key, scratch_1);
253 } else {
254 panic!("auto_key[{}] not found", gal_el);
255 }
256
257 // a = a * X^-t + b - phi(a * X^-t - b)
258 a.sub_inplace_ab(module, &tmp_b);
259 a.normalize_inplace(module, scratch_1);
260
261 // a = a + b * X^t - phi(a * X^-t - b) * X^t
262 // = a + b * X^t - phi(a * X^-t - b) * - phi(X^t)
263 // = a + b * X^t + phi(a - b * X^t)
264 a.rotate_inplace(module, t);
265 } else {
266 a.rsh(1, scratch);
267 // a = a + phi(a)
268 if let Some(key) = auto_keys.get(&gal_el) {
269 a.automorphism_add_inplace(module, key, scratch);
270 } else {
271 panic!("auto_key[{}] not found", gal_el);
272 }
273 }
274 } else {
275 if let Some(b) = b {
276 let (mut tmp_b, scratch_1) = scratch.tmp_glwe_ct(module, basek, k, rank);
277 tmp_b.rotate(module, 1 << (log_n - i - 1), b);
278 tmp_b.rsh(1, scratch_1);
279
280 // a = (b* X^t - phi(b* X^t))
281 if let Some(key) = auto_keys.get(&gal_el) {
282 a.automorphism_sub_ba::<&mut [u8], _>(module, &tmp_b, key, scratch_1);
283 } else {
284 panic!("auto_key[{}] not found", gal_el);
285 }
286
287 acc.value = true;
288 }
289 }
290}