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