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