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
 23// TODO(Jay): Shouldn't Decompose also return corresponding gadget vector ?
 24pub struct DefaultDecomposer<T> {
 25    q: T,
 26    logq: usize,
 27    logb: usize,
 28    d: usize,
 29    ignore_bits: usize,
 30    ignore_limbs: usize,
 31}
 32
 33pub trait NumInfo {
 34    const BITS: u32;
 35}
 36
 37impl NumInfo for u64 {
 38    const BITS: u32 = u64::BITS;
 39}
 40impl NumInfo for u32 {
 41    const BITS: u32 = u32::BITS;
 42}
 43impl NumInfo for u128 {
 44    const BITS: u32 = u128::BITS;
 45}
 46
 47impl<T: PrimInt + NumInfo + Debug> DefaultDecomposer<T> {
 48    pub fn new(q: T, logb: usize, d: usize) -> DefaultDecomposer<T> {
 49        // if q is power of 2, then `BITS - leading_zeros` outputs logq + 1.
 50        let logq = if q & (q - T::one()) == T::zero() {
 51            (T::BITS - q.leading_zeros() - 1) as usize
 52        } else {
 53            (T::BITS - q.leading_zeros()) as usize
 54        };
 55
 56        let d_ideal = (logq as f64 / logb as f64).ceil().to_usize().unwrap();
 57        let ignore_limbs = (d_ideal - d);
 58        let ignore_bits = (d_ideal - d) * logb;
 59
 60        DefaultDecomposer {
 61            q,
 62            logq,
 63            logb,
 64            d,
 65            ignore_bits,
 66            ignore_limbs,
 67        }
 68    }
 69
 70    fn recompose<Op>(&self, limbs: &[T], modq_op: &Op) -> T
 71    where
 72        Op: ArithmeticOps<Element = T>,
 73    {
 74        let mut value = T::zero();
 75        for i in 0..self.d {
 76            value = modq_op.add(
 77                &value,
 78                &(modq_op.mul(
 79                    &limbs[i],
 80                    &(T::one() << (self.logb * (i + self.ignore_limbs))),
 81                )),
 82            )
 83        }
 84        value
 85    }
 86}
 87
 88impl<T: PrimInt + WrappingSub + Debug> Decomposer for DefaultDecomposer<T> {
 89    type Element = T;
 90    fn decompose(&self, value: &T) -> Vec<T> {
 91        let mut value = round_value(*value, self.ignore_bits);
 92
 93        let q = self.q;
 94        // if value >= (q >> 1) {
 95        //     value = value.wrapping_sub(&q);
 96        // }
 97
 98        let logb = self.logb;
 99        let b = T::one() << logb; // base
100        let b_by2 = T::one() << (logb - 1);
101        // let neg_b_by2_modq = q - b_by2;
102        let full_mask = (T::one() << logb) - T::one();
103        // let half_mask = b_by2 - T::one();
104        let mut carry = T::zero();
105        let mut out = Vec::<T>::with_capacity(self.d);
106        for i in 0..self.d {
107            let mut limb = ((value >> (logb * i)) & full_mask) + carry;
108            carry = T::zero();
109            if limb >= b_by2 {
110                limb = (q + limb) - b;
111                carry = T::one();
112            }
113
114            // carry = ((q + g - limb) % q) >> logb;
115
116            // carry = limb & b_by2;
117            // limb = (q + limb) - (carry << 1);
118            // if limb > q {
119            //     limb = limb - q;
120            // }
121            out.push(limb);
122
123            // carry = carry >> (logb - 1);
124        }
125
126        out[self.d - 1] = out[self.d - 1] + (carry << logb);
127        if out[self.d - 1] > q {
128            out[self.d - 1] = out[self.d - 1] - q;
129        }
130
131        return out;
132    }
133
134    fn d(&self) -> usize {
135        self.d
136    }
137}
138
139fn round_value<T: PrimInt>(value: T, ignore_bits: usize) -> T {
140    if ignore_bits == 0 {
141        return value;
142    }
143
144    let ignored_msb = (value & ((T::one() << ignore_bits) - T::one())) >> (ignore_bits - 1);
145    (value >> ignore_bits) + ignored_msb
146}
147
148#[cfg(test)]
149mod tests {
150    use rand::{thread_rng, Rng};
151
152    use crate::{
153        backend::{ModInit, ModularOpsU64},
154        decomposer::round_value,
155        utils::generate_prime,
156    };
157
158    use super::{Decomposer, DefaultDecomposer};
159
160    #[test]
161    fn decomposition_works() {
162        let logq = 55;
163        let logb = 9;
164        let d = 6;
165
166        let mut rng = thread_rng();
167
168        // q is prime of bits logq and i is true, other q = 1<<logq
169        // FIXME: Test fails when q is prime, albeit the difference is minute
170        for i in [true, false] {
171            let q = if i {
172                generate_prime(logq, 1 << 4, 1u64 << logq).unwrap()
173            } else {
174                1u64 << logq
175            };
176            let decomposer = DefaultDecomposer::new(q, logb, d);
177            let modq_op = ModularOpsU64::new(q);
178            for _ in 0..1000 {
179                let value = rng.gen_range(0..q);
180                let limbs = decomposer.decompose(&value);
181                let value_back = decomposer.recompose(&limbs, &modq_op);
182                let rounded_value =
183                    round_value(value, decomposer.ignore_bits) << decomposer.ignore_bits;
184                // dbg!(&limbs, q);
185                assert_eq!(
186                    rounded_value, value_back,
187                    "Expected {rounded_value} got {value_back} for q={q}"
188                );
189            }
190        }
191    }
192}