phantom-zone/src/decomposer.rs
src/decomposer.rsRaw
1use itertools::{assert_equal, izip, Itertools};
2use num_traits::{FromPrimitive, PrimInt, ToPrimitive, WrappingAdd, WrappingSub};
3use std::fmt::{Debug, Display};
4
5use crate::{
6 backend::ArithmeticOps,
7 parameters::{
8 DecompositionCount, DecompostionLogBase, DoubleDecomposerParams, SingleDecomposerParams,
9 },
10};
11
12fn gadget_vector<T: PrimInt>(logq: usize, logb: usize, d: usize) -> Vec<T> {
13 assert!(logq >= (logb * d));
14 let ignored_bits = logq - (logb * d);
15
16 (0..d)
17 .into_iter()
18 .map(|i| T::one() << (logb * i + ignored_bits))
19 .collect_vec()
20}
21
22pub trait RlweDecomposer {
23 type Element;
24 type D: Decomposer<Element = Self::Element>;
25
26 /// Decomposer for RLWE Part A
27 fn a(&self) -> &Self::D;
28 /// Decomposer for RLWE Part B
29 fn b(&self) -> &Self::D;
30}
31
32impl<D> RlweDecomposer for (D, D)
33where
34 D: Decomposer,
35{
36 type D = D;
37 type Element = D::Element;
38 fn a(&self) -> &Self::D {
39 &self.0
40 }
41 fn b(&self) -> &Self::D {
42 &self.1
43 }
44}
45
46impl<D> DoubleDecomposerParams for D
47where
48 D: RlweDecomposer,
49{
50 type Base = DecompostionLogBase;
51 type Count = DecompositionCount;
52
53 fn decomposition_base(&self) -> Self::Base {
54 assert!(
55 Decomposer::decomposition_base(self.a()) == Decomposer::decomposition_base(self.b())
56 );
57 Decomposer::decomposition_base(self.a())
58 }
59 fn decomposition_count_a(&self) -> Self::Count {
60 Decomposer::decomposition_count(self.a())
61 }
62 fn decomposition_count_b(&self) -> Self::Count {
63 Decomposer::decomposition_count(self.b())
64 }
65}
66
67impl<D> SingleDecomposerParams for D
68where
69 D: Decomposer,
70{
71 type Base = DecompostionLogBase;
72 type Count = DecompositionCount;
73 fn decomposition_base(&self) -> Self::Base {
74 Decomposer::decomposition_base(self)
75 }
76 fn decomposition_count(&self) -> Self::Count {
77 Decomposer::decomposition_count(self)
78 }
79}
80
81pub trait Decomposer {
82 type Element;
83 type Iter: Iterator<Item = Self::Element>;
84 fn new(q: Self::Element, logb: usize, d: usize) -> Self;
85
86 fn decompose_to_vec(&self, v: &Self::Element) -> Vec<Self::Element>;
87 fn decompose_iter(&self, v: &Self::Element) -> Self::Iter;
88 fn decomposition_count(&self) -> DecompositionCount;
89 fn decomposition_base(&self) -> DecompostionLogBase;
90 fn gadget_vector(&self) -> Vec<Self::Element>;
91}
92
93pub struct DefaultDecomposer<T> {
94 /// Ciphertext modulus
95 q: T,
96 /// Log of ciphertext modulus
97 logq: usize,
98 /// Log of base B
99 logb: usize,
100 /// base B
101 b: T,
102 /// (B - 1). To simulate (% B) as &(B-1), that is extract least significant
103 /// logb bits
104 b_mask: T,
105 /// B/2
106 bby2: T,
107 /// Decomposition count
108 d: usize,
109 /// No. of bits to ignore in rounding
110 ignore_bits: usize,
111}
112
113pub trait NumInfo {
114 const BITS: u32;
115}
116
117impl NumInfo for u64 {
118 const BITS: u32 = u64::BITS;
119}
120impl NumInfo for u32 {
121 const BITS: u32 = u32::BITS;
122}
123impl NumInfo for u128 {
124 const BITS: u32 = u128::BITS;
125}
126
127impl<T: PrimInt + NumInfo + Debug> DefaultDecomposer<T> {
128 fn recompose<Op>(&self, limbs: &[T], modq_op: &Op) -> T
129 where
130 Op: ArithmeticOps<Element = T>,
131 {
132 let mut value = T::zero();
133 let gadget_vector = gadget_vector(self.logq, self.logb, self.d);
134 assert!(limbs.len() == gadget_vector.len());
135 izip!(limbs.iter(), gadget_vector.iter())
136 .for_each(|(d_el, beta)| value = modq_op.add(&value, &modq_op.mul(d_el, beta)));
137
138 value
139 }
140}
141
142impl<
143 T: PrimInt
144 + ToPrimitive
145 + FromPrimitive
146 + WrappingSub
147 + WrappingAdd
148 + NumInfo
149 + From<bool>
150 + Display
151 + Debug,
152 > Decomposer for DefaultDecomposer<T>
153{
154 type Element = T;
155 type Iter = DecomposerIter<T>;
156
157 fn new(q: T, logb: usize, d: usize) -> DefaultDecomposer<T> {
158 // if q is power of 2, then `BITS - leading_zeros` outputs logq + 1.
159 let logq = if q & (q - T::one()) == T::zero() {
160 (T::BITS - q.leading_zeros() - 1) as usize
161 } else {
162 (T::BITS - q.leading_zeros()) as usize
163 };
164
165 assert!(
166 logq >= (logb * d),
167 "Decomposer wants logq >= logb*d but got logq={logq}, logb={logb}, d={d}"
168 );
169
170 let ignore_bits = logq - (logb * d);
171
172 DefaultDecomposer {
173 q,
174 logq,
175 logb,
176 b: T::one() << logb,
177 b_mask: (T::one() << logb) - T::one(),
178 bby2: T::one() << (logb - 1),
179 d,
180 ignore_bits,
181 }
182 }
183
184 fn decompose_to_vec(&self, value: &T) -> Vec<T> {
185 let q = self.q;
186 let logb = self.logb;
187 let b = T::one() << logb;
188 let full_mask = b - T::one();
189 let bby2 = b >> 1;
190
191 let mut value = *value;
192 if value >= (q >> 1) {
193 value = !(q - value) + T::one()
194 }
195 value = round_value(value, self.ignore_bits);
196 let mut out = Vec::with_capacity(self.d);
197 for _ in 0..(self.d) {
198 let k_i = value & full_mask;
199
200 value = (value - k_i) >> logb;
201
202 if k_i > bby2 || (k_i == bby2 && ((value & T::one()) == T::one())) {
203 out.push(q - (b - k_i));
204 value = value + T::one();
205 } else {
206 out.push(k_i);
207 }
208 }
209
210 return out;
211 }
212
213 fn decomposition_count(&self) -> DecompositionCount {
214 DecompositionCount(self.d)
215 }
216
217 fn decomposition_base(&self) -> DecompostionLogBase {
218 DecompostionLogBase(self.logb)
219 }
220
221 fn decompose_iter(&self, value: &T) -> DecomposerIter<T> {
222 let mut value = *value;
223 if value >= (self.q >> 1) {
224 value = !(self.q - value) + T::one()
225 }
226 value = round_value(value, self.ignore_bits);
227
228 DecomposerIter {
229 value,
230 q: self.q,
231 logq: self.logq,
232 logb: self.logb,
233 b: self.b,
234 bby2: self.bby2,
235 b_mask: self.b_mask,
236 steps_left: self.d,
237 }
238 }
239
240 fn gadget_vector(&self) -> Vec<T> {
241 return gadget_vector(self.logq, self.logb, self.d);
242 }
243}
244
245impl<T: PrimInt> DefaultDecomposer<T> {}
246
247pub struct DecomposerIter<T> {
248 /// Value to decompose
249 value: T,
250 steps_left: usize,
251 /// (1 << logb) - 1 (for % (1<<logb); i.e. to extract least signiciant logb
252 /// bits)
253 b_mask: T,
254 logb: usize,
255 // b/2 = 1 << (logb-1)
256 bby2: T,
257 /// Ciphertext modulus
258 q: T,
259 /// Log of ciphertext modulus
260 logq: usize,
261 /// b = 1 << logb
262 b: T,
263}
264
265impl<T: PrimInt + From<bool> + WrappingSub + Display> Iterator for DecomposerIter<T> {
266 type Item = T;
267
268 fn next(&mut self) -> Option<Self::Item> {
269 if self.steps_left != 0 {
270 self.steps_left -= 1;
271 let k_i = self.value & self.b_mask;
272
273 self.value = (self.value - k_i) >> self.logb;
274
275 // if k_i > self.bby2 || (k_i == self.bby2 && ((self.value &
276 // T::one()) == T::one())) { self.value = self.value
277 // + T::one(); Some(self.q + k_i - self.b)
278 // } else {
279 // Some(k_i)
280 // }
281
282 // Following is without branching impl of the commented version above. It
283 // happens to speed up bootstrapping for `SMALL_MP_BOOL_PARAMS` (& other
284 // parameters as well but I haven't tested) by roughly 15ms.
285 // Suprisingly the improvement does not show up when I benchmark
286 // `decomposer_iter` in isolation. Putting this remark here as a
287 // future task to investiage (TODO).
288 let carry_bool =
289 k_i > self.bby2 || (k_i == self.bby2 && ((self.value & T::one()) == T::one()));
290 let carry = <T as From<bool>>::from(carry_bool);
291 let neg_carry = (T::zero().wrapping_sub(&carry));
292 self.value = self.value + carry;
293 Some((neg_carry & self.q) + k_i - (carry << self.logb))
294
295 // Some(
296 // (self.q & ((carry << self.logq) - (T::one() & carry))) + k_i
297 // - (carry << self.logb), )
298
299 // Some(k_i)
300 } else {
301 None
302 }
303 }
304}
305
306fn round_value<T: PrimInt + WrappingAdd>(value: T, ignore_bits: usize) -> T {
307 if ignore_bits == 0 {
308 return value;
309 }
310
311 let ignored_msb = (value & ((T::one() << ignore_bits) - T::one())) >> (ignore_bits - 1);
312 (value >> ignore_bits).wrapping_add(&ignored_msb)
313}
314
315#[cfg(test)]
316mod tests {
317
318 use itertools::{izip, Itertools};
319 use rand::{thread_rng, Rng};
320
321 use crate::{
322 backend::{ModInit, ModularOpsU64, Modulus},
323 decomposer::round_value,
324 utils::{generate_prime, tests::Stats},
325 };
326
327 use super::{Decomposer, DefaultDecomposer};
328
329 #[test]
330 fn decomposition_works() {
331 let ring_size = 1 << 11;
332
333 let mut rng = thread_rng();
334
335 for logq in [37, 55] {
336 let logb = 11;
337 let d = 3;
338 // let mut stats = vec![Stats::new(); d];
339
340 for i in [true, false] {
341 let q = if i {
342 generate_prime(logq, 2 * ring_size, 1u64 << logq).unwrap()
343 } else {
344 1u64 << logq
345 };
346 let decomposer = DefaultDecomposer::new(q, logb, d);
347 let modq_op = ModularOpsU64::new(q);
348 for _ in 0..1000000 {
349 let value = rng.gen_range(0..q);
350 let limbs = decomposer.decompose_to_vec(&value);
351 let limbs_from_iter = decomposer.decompose_iter(&value).collect_vec();
352 assert_eq!(limbs, limbs_from_iter);
353 let value_back = round_value(
354 decomposer.recompose(&limbs, &modq_op),
355 decomposer.ignore_bits,
356 );
357 let rounded_value = round_value(value, decomposer.ignore_bits);
358 assert!((rounded_value as i64 - value_back as i64).abs() <= 1,);
359
360 // izip!(stats.iter_mut(), limbs.iter()).for_each(|(s, l)| {
361 // s.add_more(&vec![q.map_element_to_i64(l)]);
362 // });
363 }
364 }
365
366 // stats.iter().enumerate().for_each(|(index, s)| {
367 // println!(
368 // "Limb {index} - Mean: {}, Std: {}",
369 // s.mean(),
370 // s.std_dev().abs().log2()
371 // );
372 // });
373 }
374 }
375}