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 fn new(q: Self::Element, logb: usize, d: usize) -> Self;
43 //FIXME(Jay): there's no reason why it returns a vec instead of an iterator
44 fn decompose(&self, v: &Self::Element) -> Vec<Self::Element>;
45 fn decomposition_count(&self) -> usize;
46}
47
48// TODO(Jay): Shouldn't Decompose also return corresponding gadget vector ?
49pub struct DefaultDecomposer<T> {
50 q: T,
51 logq: usize,
52 logb: usize,
53 d: usize,
54 ignore_bits: usize,
55 ignore_limbs: usize,
56}
57
58pub trait NumInfo {
59 const BITS: u32;
60}
61
62impl NumInfo for u64 {
63 const BITS: u32 = u64::BITS;
64}
65impl NumInfo for u32 {
66 const BITS: u32 = u32::BITS;
67}
68impl NumInfo for u128 {
69 const BITS: u32 = u128::BITS;
70}
71
72impl<T: PrimInt + NumInfo + Debug> DefaultDecomposer<T> {
73 fn recompose<Op>(&self, limbs: &[T], modq_op: &Op) -> T
74 where
75 Op: ArithmeticOps<Element = T>,
76 {
77 let mut value = T::zero();
78 for i in 0..self.d {
79 value = modq_op.add(
80 &value,
81 &(modq_op.mul(
82 &limbs[i],
83 &(T::one() << (self.logb * (i + self.ignore_limbs))),
84 )),
85 )
86 }
87 value
88 }
89
90 pub(crate) fn gadget_vector(&self) -> Vec<T> {
91 return gadget_vector(self.logq, self.logb, self.d);
92 }
93}
94
95impl<T: PrimInt + ToPrimitive + FromPrimitive + WrappingSub + NumInfo> Decomposer
96 for DefaultDecomposer<T>
97{
98 type Element = T;
99
100 fn new(q: T, logb: usize, d: usize) -> DefaultDecomposer<T> {
101 // if q is power of 2, then `BITS - leading_zeros` outputs logq + 1.
102 let logq = if q & (q - T::one()) == T::zero() {
103 (T::BITS - q.leading_zeros() - 1) as usize
104 } else {
105 (T::BITS - q.leading_zeros()) as usize
106 };
107
108 let d_ideal = (logq as f64 / logb as f64).ceil().to_usize().unwrap();
109 let ignore_limbs = (d_ideal - d);
110 let ignore_bits = (d_ideal - d) * logb;
111
112 DefaultDecomposer {
113 q,
114 logq,
115 logb,
116 d,
117 ignore_bits,
118 ignore_limbs,
119 }
120 }
121
122 /// Signed BNAF decomposition. Only returns most significant `d`
123 /// decomposition limbs
124 ///
125 /// Implements algorithm 3 of https://eprint.iacr.org/2021/1161.pdf
126 fn decompose(&self, value: &T) -> Vec<T> {
127 let mut value = round_value(*value, self.ignore_bits);
128
129 let q = self.q;
130 let logb = self.logb;
131 let b = T::one() << logb;
132 let full_mask = b - T::one();
133 let bby2 = b >> 1;
134
135 if value >= (q >> 1) {
136 value = !(q - value) + T::one()
137 }
138
139 let mut out = Vec::with_capacity(self.d);
140 for _ in 0..self.d {
141 let k_i = value & full_mask;
142
143 value = (value - k_i) >> logb;
144
145 if k_i > bby2 || (k_i == bby2 && ((value & T::one()) == T::one())) {
146 out.push(q - (b - k_i));
147 value = value + T::one();
148 } else {
149 out.push(k_i);
150 }
151 }
152
153 return out;
154 }
155
156 fn decomposition_count(&self) -> usize {
157 self.d
158 }
159}
160
161fn round_value<T: PrimInt>(value: T, ignore_bits: usize) -> T {
162 if ignore_bits == 0 {
163 return value;
164 }
165
166 let ignored_msb = (value & ((T::one() << ignore_bits) - T::one())) >> (ignore_bits - 1);
167 (value >> ignore_bits) + ignored_msb
168}
169
170#[cfg(test)]
171mod tests {
172 use num_traits::Float;
173 use rand::{thread_rng, Rng};
174
175 use crate::{
176 backend::{ModInit, ModularOpsU64},
177 decomposer::round_value,
178 utils::{generate_prime, Stats, TryConvertFrom1},
179 };
180
181 use super::{Decomposer, DefaultDecomposer};
182
183 #[test]
184 fn decomposition_works() {
185 let logq = 55;
186 let logb = 11;
187 let d = 5;
188 let ring_size = 1 << 11;
189
190 let mut rng = thread_rng();
191 let mut stats = Stats { samples: vec![] };
192
193 for i in [true] {
194 let q = if i {
195 generate_prime(logq, 2 * ring_size, 1u64 << logq).unwrap()
196 } else {
197 1u64 << logq
198 };
199 let decomposer = DefaultDecomposer::new(q, logb, d);
200 let modq_op = ModularOpsU64::new(q);
201 for _ in 0..100000 {
202 let value = rng.gen_range(0..q);
203 let limbs = decomposer.decompose(&value);
204 let value_back = decomposer.recompose(&limbs, &modq_op);
205 let rounded_value =
206 round_value(value, decomposer.ignore_bits) << decomposer.ignore_bits;
207 stats.add_more(&Vec::<i64>::try_convert_from(&limbs, &q));
208 assert_eq!(
209 rounded_value, value_back,
210 "Expected {rounded_value} got {value_back} for q={q}"
211 );
212 }
213 }
214 println!("Mean: {}", stats.mean());
215 println!("Std: {}", stats.std_dev().abs().log2());
216 }
217}