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