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