1use std::collections::HashMap;
  2
  3use poulpy_hal::{
  4    api::ModuleLogN,
  5    layouts::{Backend, GaloisElement, Module, Scratch},
  6};
  7
  8use crate::{
  9    GLWEAdd, GLWEAutomorphism, GLWECopy, GLWENormalize, GLWERotate, GLWEShift, GLWESub, ScratchTakeCore,
 10    layouts::{GGLWEInfos, GGLWEPreparedToRef, GLWEInfos, GLWEToMut, GLWEToRef, GetGaloisElement},
 11};
 12pub trait GLWEPacking<BE: Backend> {
 13    /// Packs [x_0: GLWE(m_0), x_1: GLWE(m_1), ..., x_i: GLWE(m_i)]
 14    /// to [0: GLWE(m_0 * X^x_0 + m_1 * X^x_1 + ... + m_i * X^x_i)]
 15    fn glwe_pack<R, K>(
 16        &self,
 17        cts: &mut HashMap<usize, &mut R>,
 18        log_gap_out: usize,
 19        keys: &HashMap<i64, K>,
 20        scratch: &mut Scratch<BE>,
 21    ) where
 22        R: GLWEToMut + GLWEToRef + GLWEInfos,
 23        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos;
 24}
 25
 26impl<BE: Backend> GLWEPacking<BE> for Module<BE>
 27where
 28    Self: GLWEAutomorphism<BE>
 29        + GaloisElement
 30        + ModuleLogN
 31        + GLWERotate<BE>
 32        + GLWESub
 33        + GLWEShift<BE>
 34        + GLWEAdd
 35        + GLWENormalize<BE>
 36        + GLWECopy,
 37    Scratch<BE>: ScratchTakeCore<BE>,
 38{
 39    /// Packs [x_0: GLWE(m_0), x_1: GLWE(m_1), ..., x_i: GLWE(m_i)]
 40    /// to [0: GLWE(m_0 * X^x_0 + m_1 * X^x_1 + ... + m_i * X^x_i)]
 41    fn glwe_pack<R, K>(
 42        &self,
 43        cts: &mut HashMap<usize, &mut R>,
 44        log_gap_out: usize,
 45        keys: &HashMap<i64, K>,
 46        scratch: &mut Scratch<BE>,
 47    ) where
 48        R: GLWEToMut + GLWEToRef + GLWEInfos,
 49        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
 50    {
 51        #[cfg(debug_assertions)]
 52        {
 53            assert!(*cts.keys().max().unwrap() < self.n())
 54        }
 55
 56        let log_n: usize = self.log_n();
 57
 58        for i in 0..(log_n - log_gap_out) {
 59            let t: usize = (1 << log_n).min(1 << (log_n - 1 - i));
 60
 61            let key: &K = if i == 0 {
 62                keys.get(&-1).unwrap()
 63            } else {
 64                keys.get(&self.galois_element(1 << (i - 1))).unwrap()
 65            };
 66
 67            for j in 0..t {
 68                let mut a: Option<&mut R> = cts.remove(&j);
 69                let mut b: Option<&mut R> = cts.remove(&(j + t));
 70
 71                pack_internal(self, &mut a, &mut b, i, key, scratch);
 72
 73                if let Some(a) = a {
 74                    cts.insert(j, a);
 75                } else if let Some(b) = b {
 76                    cts.insert(j, b);
 77                }
 78            }
 79        }
 80    }
 81}
 82
 83#[allow(clippy::too_many_arguments)]
 84fn pack_internal<M, A, B, K, BE: Backend>(
 85    module: &M,
 86    a: &mut Option<&mut A>,
 87    b: &mut Option<&mut B>,
 88    i: usize,
 89    auto_key: &K,
 90    scratch: &mut Scratch<BE>,
 91) where
 92    M: GLWEAutomorphism<BE> + GLWERotate<BE> + GLWESub + GLWEShift<BE> + GLWEAdd + GLWENormalize<BE>,
 93    A: GLWEToMut + GLWEToRef + GLWEInfos,
 94    B: GLWEToMut + GLWEToRef + GLWEInfos,
 95    K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
 96    Scratch<BE>: ScratchTakeCore<BE>,
 97{
 98    // Goal is to evaluate: a = a + b*X^t + phi(a - b*X^t))
 99    // We also use the identity: AUTO(a * X^t, g) = -X^t * AUTO(a, g)
100    // where t = 2^(log_n - i - 1) and g = 5^{2^(i - 1)}
101    // Different cases for wether a and/or b are zero.
102    //
103    // Implicite RSH without modulus switch, introduces extra I(X) * Q/2 on decryption.
104    // Necessary so that the scaling of the plaintext remains constant.
105    // It however is ok to do so here because coefficients are eventually
106    // either mapped to garbage or twice their value which vanishes I(X)
107    // since 2*(I(X) * Q/2) = I(X) * Q = 0 mod Q.
108    if let Some(a) = a.as_deref_mut() {
109        let t: i64 = 1 << (a.n().log2() - i - 1);
110
111        if let Some(b) = b.as_deref_mut() {
112            let (mut tmp_b, scratch_1) = scratch.take_glwe(a);
113
114            // a = a * X^-t
115            module.glwe_rotate_inplace(-t, a, scratch_1);
116
117            // tmp_b = a * X^-t - b
118            module.glwe_sub(&mut tmp_b, a, b);
119            module.glwe_rsh(1, &mut tmp_b, scratch_1);
120
121            // a = a * X^-t + b
122            module.glwe_add_inplace(a, b);
123            module.glwe_rsh(1, a, scratch_1);
124
125            module.glwe_normalize_inplace(&mut tmp_b, scratch_1);
126
127            // tmp_b = phi(a * X^-t - b)
128            module.glwe_automorphism_inplace(&mut tmp_b, auto_key, scratch_1);
129
130            // a = a * X^-t + b - phi(a * X^-t - b)
131            module.glwe_sub_inplace(a, &tmp_b);
132            module.glwe_normalize_inplace(a, scratch_1);
133
134            // a = a + b * X^t - phi(a * X^-t - b) * X^t
135            //   = a + b * X^t - phi(a * X^-t - b) * - phi(X^t)
136            //   = a + b * X^t + phi(a - b * X^t)
137            module.glwe_rotate_inplace(t, a, scratch_1);
138        } else {
139            module.glwe_rsh(1, a, scratch);
140            // a = a + phi(a)
141            module.glwe_automorphism_add_inplace(a, auto_key, scratch);
142        }
143    } else if let Some(b) = b.as_deref_mut() {
144        let t: i64 = 1 << (b.n().log2() - i - 1);
145
146        let (mut tmp_b, scratch_1) = scratch.take_glwe(b);
147        module.glwe_rotate(t, &mut tmp_b, b);
148        module.glwe_rsh(1, &mut tmp_b, scratch_1);
149
150        // a = (b* X^t - phi(b* X^t))
151        module.glwe_automorphism_sub_negate(b, &tmp_b, auto_key, scratch_1);
152    }
153}