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