poulpy_core/encryption/
glwe_tensor_key.rs

1use poulpy_hal::{
2    api::ModuleN,
3    layouts::{Backend, DataMut, Module, Scratch},
4    source::Source,
5};
6
7use crate::{
8    GGLWEEncryptSk, GetDistribution, ScratchTakeCore,
9    layouts::{
10        GGLWEInfos, GGLWELayout, GGLWEToMut, GLWEInfos, GLWESecretTensor, GLWESecretTensorFactory, GLWESecretToRef,
11        GLWETensorKey,
12        prepared::{GLWESecretPrepared, GLWESecretPreparedFactory},
13    },
14};
15
16impl GLWETensorKey<Vec<u8>> {
17    pub fn encrypt_sk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
18    where
19        A: GGLWEInfos,
20        M: GLWETensorKeyEncryptSk<BE>,
21    {
22        module.glwe_tensor_key_encrypt_sk_tmp_bytes(infos)
23    }
24}
25
26impl<DataSelf: DataMut> GLWETensorKey<DataSelf> {
27    pub fn encrypt_sk<M, S, BE: Backend>(
28        &mut self,
29        module: &M,
30        sk: &S,
31        source_xa: &mut Source,
32        source_xe: &mut Source,
33        scratch: &mut Scratch<BE>,
34    ) where
35        M: GLWETensorKeyEncryptSk<BE>,
36        S: GLWESecretToRef + GetDistribution + GLWEInfos,
37        Scratch<BE>: ScratchTakeCore<BE>,
38    {
39        module.glwe_tensor_key_encrypt_sk(self, sk, source_xa, source_xe, scratch);
40    }
41}
42
43pub trait GLWETensorKeyEncryptSk<BE: Backend> {
44    fn glwe_tensor_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
45    where
46        A: GGLWEInfos;
47
48    fn glwe_tensor_key_encrypt_sk<R, S>(
49        &self,
50        res: &mut R,
51        sk: &S,
52        source_xa: &mut Source,
53        source_xe: &mut Source,
54        scratch: &mut Scratch<BE>,
55    ) where
56        R: GGLWEToMut + GGLWEInfos,
57        S: GLWESecretToRef + GetDistribution + GLWEInfos;
58}
59
60impl<BE: Backend> GLWETensorKeyEncryptSk<BE> for Module<BE>
61where
62    Self: ModuleN + GGLWEEncryptSk<BE> + GLWESecretPreparedFactory<BE> + GLWESecretTensorFactory<BE>,
63    Scratch<BE>: ScratchTakeCore<BE>,
64{
65    fn glwe_tensor_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
66    where
67        A: GGLWEInfos,
68    {
69        let sk_prepared: usize = GLWESecretPrepared::bytes_of(self, infos.rank_out());
70        let sk_tensor: usize = GLWESecretTensor::bytes_of_from_infos(infos);
71
72        let tensor_infos: GGLWELayout = GGLWELayout {
73            n: infos.n(),
74            base2k: infos.base2k(),
75            k: infos.k(),
76            rank_in: GLWESecretTensor::pairs(infos.rank().into()).into(),
77            rank_out: infos.rank_out(),
78            dnum: infos.dnum(),
79            dsize: infos.dsize(),
80        };
81
82        let gglwe_encrypt: usize = self.gglwe_encrypt_sk_tmp_bytes(&tensor_infos);
83
84        (sk_prepared + sk_tensor) + gglwe_encrypt.max(self.glwe_secret_tensor_prepare_tmp_bytes(infos.rank()))
85    }
86
87    fn glwe_tensor_key_encrypt_sk<R, S>(
88        &self,
89        res: &mut R,
90        sk: &S,
91        source_xa: &mut Source,
92        source_xe: &mut Source,
93        scratch: &mut Scratch<BE>,
94    ) where
95        R: GGLWEToMut + GGLWEInfos,
96        S: GLWESecretToRef + GetDistribution + GLWEInfos,
97    {
98        assert_eq!(res.rank_out(), sk.rank());
99        assert_eq!(res.n(), sk.n());
100
101        let (mut sk_prepared, scratch_1) = scratch.take_glwe_secret_prepared(self, res.rank());
102        let (mut sk_tensor, scratch_2) = scratch_1.take_glwe_secret_tensor(self.n().into(), res.rank());
103        sk_prepared.prepare(self, sk);
104        sk_tensor.prepare(self, sk, scratch_2);
105
106        self.gglwe_encrypt_sk(
107            res,
108            &sk_tensor.data,
109            &sk_prepared,
110            source_xa,
111            source_xe,
112            scratch_2,
113        );
114    }
115}