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