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    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
 46pub trait Decomposer {
 47    type Element;
 48    type Iter: Iterator<Item = Self::Element>;
 49    fn new(q: Self::Element, logb: usize, d: usize) -> Self;
 50
 51    fn decompose_to_vec(&self, v: &Self::Element) -> Vec<Self::Element>;
 52    fn decompose_iter(&self, v: &Self::Element) -> Self::Iter;
 53    fn decomposition_count(&self) -> usize;
 54}
 55
 56pub struct DefaultDecomposer<T> {
 57    /// Ciphertext modulus
 58    q: T,
 59    /// Log of ciphertext modulus
 60    logq: usize,
 61    /// Log of base B
 62    logb: usize,
 63    /// base B
 64    b: T,
 65    /// (B - 1). To simulate (% B) as &(B-1), that is extract least significant
 66    /// logb bits
 67    b_mask: T,
 68    /// B/2
 69    bby2: T,
 70    /// Decomposition count
 71    d: usize,
 72    /// No. of bits to ignore in rounding
 73    ignore_bits: usize,
 74}
 75
 76pub trait NumInfo {
 77    const BITS: u32;
 78}
 79
 80impl NumInfo for u64 {
 81    const BITS: u32 = u64::BITS;
 82}
 83impl NumInfo for u32 {
 84    const BITS: u32 = u32::BITS;
 85}
 86impl NumInfo for u128 {
 87    const BITS: u32 = u128::BITS;
 88}
 89
 90impl<T: PrimInt + NumInfo + Debug> DefaultDecomposer<T> {
 91    fn recompose<Op>(&self, limbs: &[T], modq_op: &Op) -> T
 92    where
 93        Op: ArithmeticOps<Element = T>,
 94    {
 95        let mut value = T::zero();
 96        let gadget_vector = self.gadget_vector();
 97        assert!(limbs.len() == gadget_vector.len());
 98        izip!(limbs.iter(), gadget_vector.iter())
 99            .for_each(|(d_el, beta)| value = modq_op.add(&value, &modq_op.mul(d_el, beta)));
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<
110        T: PrimInt
111            + ToPrimitive
112            + FromPrimitive
113            + WrappingSub
114            + WrappingAdd
115            + NumInfo
116            + From<bool>
117            + Display,
118    > Decomposer for DefaultDecomposer<T>
119{
120    type Element = T;
121    type Iter = DecomposerIter<T>;
122
123    fn new(q: T, logb: usize, d: usize) -> DefaultDecomposer<T> {
124        // if q is power of 2, then `BITS - leading_zeros` outputs logq + 1.
125        let logq = if q & (q - T::one()) == T::zero() {
126            (T::BITS - q.leading_zeros() - 1) as usize
127        } else {
128            (T::BITS - q.leading_zeros()) as usize
129        };
130
131        let ignore_bits = logq - (logb * d);
132
133        DefaultDecomposer {
134            q,
135            logq,
136            logb,
137            b: T::one() << logb,
138            b_mask: (T::one() << logb) - T::one(),
139            bby2: T::one() << (logb - 1),
140            d,
141            ignore_bits,
142        }
143    }
144
145    // TODO(Jay): Outline the caveat
146    fn decompose_to_vec(&self, value: &T) -> Vec<T> {
147        let mut value = round_value(*value, self.ignore_bits);
148
149        let q = self.q;
150        let logb = self.logb;
151        let b = T::one() << logb;
152        let full_mask = b - T::one();
153        let bby2 = b >> 1;
154
155        if value >= (q >> 1) {
156            value = !(q - value) + T::one()
157        }
158
159        let mut out = Vec::with_capacity(self.d);
160        for _ in 0..self.d {
161            let k_i = value & full_mask;
162
163            value = (value - k_i) >> logb;
164
165            if k_i > bby2 || (k_i == bby2 && ((value & T::one()) == T::one())) {
166                out.push(q - (b - k_i));
167                value = value + T::one();
168            } else {
169                out.push(k_i);
170            }
171        }
172
173        return out;
174    }
175
176    fn decomposition_count(&self) -> usize {
177        self.d
178    }
179
180    fn decompose_iter(&self, value: &T) -> DecomposerIter<T> {
181        let mut value = round_value(*value, self.ignore_bits);
182
183        if value >= (self.q >> 1) {
184            value = !(self.q - value) + T::one()
185        }
186
187        DecomposerIter {
188            value,
189            q: self.q,
190            logq: self.logq,
191            logb: self.logb,
192            b: self.b,
193            bby2: self.bby2,
194            b_mask: self.b_mask,
195            steps_left: self.d,
196        }
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, TryConvertFrom1},
280    };
281
282    use super::{Decomposer, DefaultDecomposer};
283
284    #[test]
285    fn decomposition_works() {
286        let logq = 55;
287        let logb = 12;
288        let d = 4;
289        let ring_size = 1 << 11;
290
291        let mut rng = thread_rng();
292        let mut stats = vec![Stats::new(); d];
293
294        for i in [true] {
295            let q = if i {
296                generate_prime(logq, 2 * ring_size, 1u64 << logq).unwrap()
297            } else {
298                1u64 << logq
299            };
300            let decomposer = DefaultDecomposer::new(q, logb, d);
301            dbg!(decomposer.ignore_bits);
302            let modq_op = ModularOpsU64::new(q);
303            for _ in 0..100000 {
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_eq!(
314                //     rounded_value, value_back,
315                //     "Expected {rounded_value} got {value_back} for q={q}"
316                // );
317
318                izip!(stats.iter_mut(), limbs.iter()).for_each(|(s, l)| {
319                    s.add_more(&vec![q.map_element_to_i64(l)]);
320                });
321            }
322        }
323
324        stats.iter().enumerate().for_each(|(index, s)| {
325            println!(
326                "Limb {index} - Mean: {}, Std: {}",
327                s.mean(),
328                s.std_dev().abs().log2()
329            );
330        });
331    }
332}