1use itertools::Itertools;
  2use num_traits::{AsPrimitive, FromPrimitive, 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    type Iter: Iterator<Item = Self::Element>;
 43    fn new(q: Self::Element, logb: usize, d: usize) -> Self;
 44
 45    fn decompose_to_vec(&self, v: &Self::Element) -> Vec<Self::Element>;
 46    fn decompose_iter(&self, v: &Self::Element) -> Self::Iter;
 47    fn decomposition_count(&self) -> usize;
 48}
 49
 50pub struct DefaultDecomposer<T> {
 51    /// Ciphertext modulus
 52    q: T,
 53    /// Log of ciphertext modulus
 54    logq: usize,
 55    /// Log of base B
 56    logb: usize,
 57    /// base B
 58    b: T,
 59    /// (B - 1). To simulate (% B) as &(B-1), that is extract least significant
 60    /// logb bits
 61    b_mask: T,
 62    /// B/2
 63    bby2: T,
 64    /// Decomposition count
 65    d: usize,
 66    /// No. of bits to ignore in rounding
 67    ignore_bits: usize,
 68    /// No. of limbs to ignore in rounding. Set to ceil(logq / logb) - d
 69    ignore_limbs: usize,
 70}
 71
 72pub trait NumInfo {
 73    const BITS: u32;
 74}
 75
 76impl NumInfo for u64 {
 77    const BITS: u32 = u64::BITS;
 78}
 79impl NumInfo for u32 {
 80    const BITS: u32 = u32::BITS;
 81}
 82impl NumInfo for u128 {
 83    const BITS: u32 = u128::BITS;
 84}
 85
 86impl<T: PrimInt + NumInfo + Debug> DefaultDecomposer<T> {
 87    fn recompose<Op>(&self, limbs: &[T], modq_op: &Op) -> T
 88    where
 89        Op: ArithmeticOps<Element = T>,
 90    {
 91        let mut value = T::zero();
 92        for i in 0..self.d {
 93            value = modq_op.add(
 94                &value,
 95                &(modq_op.mul(
 96                    &limbs[i],
 97                    &(T::one() << (self.logb * (i + self.ignore_limbs))),
 98                )),
 99            )
100        }
101        value
102    }
103
104    pub(crate) fn gadget_vector(&self) -> Vec<T> {
105        return gadget_vector(self.logq, self.logb, self.d);
106    }
107}
108
109impl<T: PrimInt + ToPrimitive + FromPrimitive + WrappingSub + NumInfo> Decomposer
110    for DefaultDecomposer<T>
111{
112    type Element = T;
113    type Iter = DecomposerIter<T>;
114
115    fn new(q: T, logb: usize, d: usize) -> DefaultDecomposer<T> {
116        // if q is power of 2, then `BITS - leading_zeros` outputs logq + 1.
117        let logq = if q & (q - T::one()) == T::zero() {
118            (T::BITS - q.leading_zeros() - 1) as usize
119        } else {
120            (T::BITS - q.leading_zeros()) as usize
121        };
122
123        let d_ideal = (logq as f64 / logb as f64).ceil().to_usize().unwrap();
124        let ignore_limbs = (d_ideal - d);
125        let ignore_bits = (d_ideal - d) * logb;
126
127        DefaultDecomposer {
128            q,
129            logq,
130            logb,
131            b: T::one() << logb,
132            b_mask: (T::one() << logb) - T::one(),
133            bby2: T::one() << (logb - 1),
134            d,
135            ignore_bits,
136            ignore_limbs,
137        }
138    }
139
140    // TODO(Jay): Outline the caveat
141    fn decompose_to_vec(&self, value: &T) -> Vec<T> {
142        let mut value = round_value(*value, self.ignore_bits);
143
144        let q = self.q;
145        let logb = self.logb;
146        let b = T::one() << logb;
147        let full_mask = b - T::one();
148        let bby2 = b >> 1;
149
150        if value >= (q >> 1) {
151            value = !(q - value) + T::one()
152        }
153
154        let mut out = Vec::with_capacity(self.d);
155        for _ in 0..self.d {
156            let k_i = value & full_mask;
157
158            value = (value - k_i) >> logb;
159
160            if k_i > bby2 || (k_i == bby2 && ((value & T::one()) == T::one())) {
161                out.push(q - (b - k_i));
162                value = value + T::one();
163            } else {
164                out.push(k_i);
165            }
166        }
167
168        return out;
169    }
170
171    fn decomposition_count(&self) -> usize {
172        self.d
173    }
174
175    fn decompose_iter(&self, value: &T) -> DecomposerIter<T> {
176        let mut value = round_value(*value, self.ignore_bits);
177
178        if value >= (self.q >> 1) {
179            value = !(self.q - value) + T::one()
180        }
181
182        DecomposerIter {
183            value,
184            q: self.q,
185            logb: self.logb,
186            b: self.b,
187            bby2: self.bby2,
188            b_mask: self.b_mask,
189            steps_left: self.d,
190        }
191    }
192}
193
194impl<T: PrimInt> DefaultDecomposer<T> {}
195
196pub struct DecomposerIter<T> {
197    /// Value to decompose
198    value: T,
199    steps_left: usize,
200    /// (1 << logb) - 1 (for % (1<<logb); i.e. to extract least signiciant logb
201    /// bits)
202    b_mask: T,
203    logb: usize,
204    // b/2 = 1 << (logb-1)
205    bby2: T,
206    /// Ciphertext modulus
207    q: T,
208    /// b = 1 << logb
209    b: T,
210}
211
212impl<T: PrimInt> Iterator for DecomposerIter<T> {
213    type Item = T;
214
215    fn next(&mut self) -> Option<Self::Item> {
216        if self.steps_left != 0 {
217            self.steps_left -= 1;
218            let k_i = self.value & self.b_mask;
219
220            self.value = (self.value - k_i) >> self.logb;
221
222            if k_i > self.bby2 || (k_i == self.bby2 && ((self.value & T::one()) == T::one())) {
223                self.value = self.value + T::one();
224                Some(self.q + k_i - self.b)
225            } else {
226                Some(k_i)
227            }
228
229            // let carry = <T as From<bool>>::from(
230            //     k_i > self.bby2 || (k_i == self.bby2 && ((self.value &
231            // T::one()) == T::one())), );
232            // self.value = self.value + carry;
233
234            // Some(
235            //     (self.q & ((carry << self.logq) - (T::one() & carry))) + k_i
236            // - (carry << self.logb), )
237
238            // Some(k_i)
239        } else {
240            None
241        }
242    }
243}
244
245fn round_value<T: PrimInt>(value: T, ignore_bits: usize) -> T {
246    if ignore_bits == 0 {
247        return value;
248    }
249
250    let ignored_msb = (value & ((T::one() << ignore_bits) - T::one())) >> (ignore_bits - 1);
251    (value >> ignore_bits) + ignored_msb
252}
253
254#[cfg(test)]
255mod tests {
256
257    use rand::{thread_rng, Rng};
258
259    use crate::{
260        backend::{ModInit, ModularOpsU64},
261        decomposer::round_value,
262        utils::{generate_prime, Stats, TryConvertFrom1},
263    };
264
265    use super::{Decomposer, DefaultDecomposer};
266
267    #[test]
268    fn decomposition_works() {
269        let logq = 55;
270        let logb = 11;
271        let d = 5;
272        let ring_size = 1 << 11;
273
274        let mut rng = thread_rng();
275        let mut stats = Stats { samples: vec![] };
276
277        for i in [true] {
278            let q = if i {
279                generate_prime(logq, 2 * ring_size, 1u64 << logq).unwrap()
280            } else {
281                1u64 << logq
282            };
283            let decomposer = DefaultDecomposer::new(q, logb, d);
284            let modq_op = ModularOpsU64::new(q);
285            for _ in 0..100000 {
286                let value = rng.gen_range(0..q);
287                let limbs = decomposer.decompose_to_vec(&value);
288                let value_back = round_value(
289                    decomposer.recompose(&limbs, &modq_op),
290                    decomposer.ignore_bits,
291                );
292                let rounded_value = round_value(value, decomposer.ignore_bits);
293                assert_eq!(
294                    rounded_value, value_back,
295                    "Expected {rounded_value} got {value_back} for q={q}"
296                );
297
298                stats.add_more(&Vec::<i64>::try_convert_from(&limbs, &q));
299            }
300        }
301        println!("Mean: {}", stats.mean());
302        println!("Std: {}", stats.std_dev().abs().log2());
303    }
304}