1use itertools::Itertools;
  2use num_traits::{AsPrimitive, One, PrimInt, ToPrimitive, WrappingSub, Zero};
  3use std::{fmt::Debug, marker::PhantomData, ops::Rem};
  4
  5use crate::backend::{ArithmeticOps, ModularOpsU64};
  6
  7pub fn 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 Decomposer {
 17    type Element;
 18    //FIXME(Jay): there's no reason why it returns a vec instead of an iterator
 19    fn decompose(&self, v: &Self::Element) -> Vec<Self::Element>;
 20    fn d(&self) -> usize;
 21}
 22
 23pub struct DefaultDecomposer<T> {
 24    q: T,
 25    logq: usize,
 26    logb: usize,
 27    d: usize,
 28    ignore_bits: usize,
 29    ignore_limbs: usize,
 30}
 31
 32pub trait NumInfo {
 33    const BITS: u32;
 34}
 35
 36impl NumInfo for u64 {
 37    const BITS: u32 = u64::BITS;
 38}
 39impl NumInfo for u32 {
 40    const BITS: u32 = u32::BITS;
 41}
 42impl NumInfo for u128 {
 43    const BITS: u32 = u128::BITS;
 44}
 45
 46impl<T: PrimInt + NumInfo + Debug> DefaultDecomposer<T> {
 47    pub fn new(q: T, logb: usize, d: usize) -> DefaultDecomposer<T> {
 48        // if q is power of 2, then BITS - leading zeros outputs logq + 1.
 49        let logq = if q & (q - T::one()) == T::zero() {
 50            (T::BITS - q.leading_zeros() - 1) as usize
 51        } else {
 52            (T::BITS - q.leading_zeros()) as usize
 53        };
 54
 55        let d_ideal = (logq as f64 / logb as f64).ceil().to_usize().unwrap();
 56        let ignore_limbs = (d_ideal - d);
 57        let ignore_bits = (d_ideal - d) * logb;
 58
 59        DefaultDecomposer {
 60            q,
 61            logq,
 62            logb,
 63            d,
 64            ignore_bits,
 65            ignore_limbs,
 66        }
 67    }
 68
 69    fn recompose<Op>(&self, limbs: &[T], modq_op: &Op) -> T
 70    where
 71        Op: ArithmeticOps<Element = T>,
 72    {
 73        let mut value = T::zero();
 74        for i in self.ignore_limbs..self.ignore_limbs + self.d {
 75            value = modq_op.add(
 76                &value,
 77                &(modq_op.mul(&limbs[i], &(T::one() << (self.logb * i)))),
 78            )
 79        }
 80        value
 81    }
 82}
 83
 84impl<T: PrimInt + WrappingSub + Debug> Decomposer for DefaultDecomposer<T> {
 85    type Element = T;
 86    fn decompose(&self, value: &T) -> Vec<T> {
 87        let value = round_value(*value, self.ignore_bits);
 88
 89        let q = self.q;
 90        let logb = self.logb;
 91        // let b = T::one() << logb; // base
 92        let b_by2 = T::one() << (logb - 1);
 93        // let neg_b_by2_modq = q - b_by2;
 94        let full_mask = (T::one() << logb) - T::one();
 95        // let half_mask = b_by2 - T::one();
 96        let mut carry = T::zero();
 97        let mut out = Vec::<T>::with_capacity(self.d);
 98        for i in 0..self.d {
 99            let mut limb = ((value >> (logb * i)) & full_mask) + carry;
100
101            carry = limb & b_by2;
102            limb = (q + limb) - (carry << 1);
103            if limb > q {
104                limb = limb - q;
105            }
106            out.push(limb);
107
108            carry = carry >> (logb - 1);
109        }
110
111        return out;
112    }
113
114    fn d(&self) -> usize {
115        self.d
116    }
117}
118
119fn round_value<T: PrimInt>(value: T, ignore_bits: usize) -> T {
120    if ignore_bits == 0 {
121        return value;
122    }
123
124    let ignored_msb = (value & ((T::one() << ignore_bits) - T::one())) >> (ignore_bits - 1);
125    (value >> ignore_bits) + ignored_msb
126}
127
128#[cfg(test)]
129mod tests {
130    use rand::{thread_rng, Rng};
131
132    use crate::{backend::ModularOpsU64, decomposer::round_value, utils::generate_prime};
133
134    use super::{Decomposer, DefaultDecomposer};
135
136    #[test]
137    fn decomposition_works() {
138        let logq = 50;
139        let logb = 5;
140        let d = 10;
141
142        // q is prime of bits logq and i is true, other q = 1<<logq
143        for i in [true, false] {
144            let q = if i {
145                generate_prime(logq, 1 << 4, 1u64 << logq).unwrap()
146            } else {
147                1u64 << 50
148            };
149
150            let decomposer = DefaultDecomposer::new(q, logb, d);
151            let modq_op = ModularOpsU64::new(q);
152            for _ in 0..100 {
153                let value = 1000000;
154                let limbs = decomposer.decompose(&value);
155                let value_back = decomposer.recompose(&limbs, &modq_op);
156                let rounded_value = round_value(value, decomposer.ignore_bits);
157                assert_eq!(
158                    rounded_value, value_back,
159                    "Expected {rounded_value} got {value_back} for q={q}"
160                );
161            }
162        }
163    }
164}