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