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