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