1use poulpy_hal::{
  2    api::{VecZnxAddNormal, VecZnxFillUniform, VecZnxNormalizeInplace},
  3    layouts::{Backend, DataMut, Module, Scratch, VecZnx, ZnxView, ZnxViewMut},
  4    source::Source,
  5};
  6
  7use crate::{
  8    ScratchTakeCore,
  9    encryption::{SIGMA, SIGMA_BOUND},
 10    layouts::{LWE, LWEInfos, LWEPlaintext, LWEPlaintextToRef, LWESecret, LWESecretToRef, LWEToMut},
 11};
 12
 13impl<DataSelf: DataMut> LWE<DataSelf> {
 14    pub fn encrypt_sk<P, S, M, BE: Backend>(
 15        &mut self,
 16        module: &M,
 17        pt: &P,
 18        sk: &S,
 19        source_xa: &mut Source,
 20        source_xe: &mut Source,
 21        scratch: &mut Scratch<BE>,
 22    ) where
 23        P: LWEPlaintextToRef,
 24        S: LWESecretToRef,
 25        M: LWEEncryptSk<BE>,
 26        Scratch<BE>: ScratchTakeCore<BE>,
 27    {
 28        module.lwe_encrypt_sk(self, pt, sk, source_xa, source_xe, scratch);
 29    }
 30}
 31
 32pub trait LWEEncryptSk<BE: Backend> {
 33    fn lwe_encrypt_sk<R, P, S>(
 34        &self,
 35        res: &mut R,
 36        pt: &P,
 37        sk: &S,
 38        source_xa: &mut Source,
 39        source_xe: &mut Source,
 40        scratch: &mut Scratch<BE>,
 41    ) where
 42        R: LWEToMut,
 43        P: LWEPlaintextToRef,
 44        S: LWESecretToRef,
 45        Scratch<BE>: ScratchTakeCore<BE>;
 46}
 47
 48impl<BE: Backend> LWEEncryptSk<BE> for Module<BE>
 49where
 50    Self: Sized + VecZnxFillUniform + VecZnxAddNormal + VecZnxNormalizeInplace<BE>,
 51{
 52    fn lwe_encrypt_sk<R, P, S>(
 53        &self,
 54        res: &mut R,
 55        pt: &P,
 56        sk: &S,
 57        source_xa: &mut Source,
 58        source_xe: &mut Source,
 59        scratch: &mut Scratch<BE>,
 60    ) where
 61        R: LWEToMut,
 62        P: LWEPlaintextToRef,
 63        S: LWESecretToRef,
 64        Scratch<BE>: ScratchTakeCore<BE>,
 65    {
 66        let res: &mut LWE<&mut [u8]> = &mut res.to_mut();
 67        let pt: &LWEPlaintext<&[u8]> = &pt.to_ref();
 68        let sk: &LWESecret<&[u8]> = &sk.to_ref();
 69
 70        #[cfg(debug_assertions)]
 71        {
 72            assert_eq!(res.n(), sk.n())
 73        }
 74
 75        let base2k: usize = res.base2k().into();
 76        let k: usize = res.k().into();
 77
 78        self.vec_znx_fill_uniform(base2k, &mut res.data, 0, source_xa);
 79
 80        let mut tmp_znx: VecZnx<Vec<u8>> = VecZnx::alloc(1, 1, res.size());
 81
 82        let min_size: usize = res.size().min(pt.size());
 83
 84        (0..min_size).for_each(|i| {
 85            tmp_znx.at_mut(0, i)[0] = pt.data.at(0, i)[0]
 86                - res.data.at(0, i)[1..]
 87                    .iter()
 88                    .zip(sk.data.at(0, 0))
 89                    .map(|(x, y)| x * y)
 90                    .sum::<i64>();
 91        });
 92
 93        (min_size..res.size()).for_each(|i| {
 94            tmp_znx.at_mut(0, i)[0] -= res.data.at(0, i)[1..]
 95                .iter()
 96                .zip(sk.data.at(0, 0))
 97                .map(|(x, y)| x * y)
 98                .sum::<i64>();
 99        });
100
101        self.vec_znx_add_normal(base2k, &mut tmp_znx, 0, k, source_xe, SIGMA, SIGMA_BOUND);
102
103        self.vec_znx_normalize_inplace(base2k, &mut tmp_znx, 0, scratch);
104
105        (0..res.size()).for_each(|i| {
106            res.data.at_mut(0, i)[0] = tmp_znx.at(0, i)[0];
107        });
108    }
109}