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