1use poulpy_hal::{
2 api::{ModuleN, ScratchAvailable, ScratchTakeBasic, SvpPrepare, VecZnxSwitchRing},
3 layouts::{Backend, DataMut, Module, ScalarZnx, Scratch},
4 source::Source,
5};
6
7use crate::{
8 ScratchTakeCore,
9 encryption::gglwe::GGLWEEncryptSk,
10 layouts::{
11 GGLWEInfos, GGLWEToMut, GLWEInfos, GLWESecret, GLWESecretToRef, GLWESwitchingKey, GLWESwitchingKeyDegreesMut, LWEInfos,
12 prepared::GLWESecretPreparedFactory,
13 },
14};
15
16impl GLWESwitchingKey<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: GLWESwitchingKeyEncryptSk<BE>,
21 {
22 module.glwe_switching_key_encrypt_sk_tmp_bytes(infos)
23 }
24
25 pub fn encrypt_pk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
26 where
27 A: GGLWEInfos,
28 M: GLWESwitchingKeyEncryptPk<BE>,
29 {
30 module.glwe_switching_key_encrypt_pk_tmp_bytes(infos)
31 }
32}
33
34impl<DataSelf: DataMut> GLWESwitchingKey<DataSelf> {
35 pub fn encrypt_sk<M, S1, S2, BE: Backend>(
36 &mut self,
37 module: &M,
38 sk_in: &S1,
39 sk_out: &S2,
40 source_xa: &mut Source,
41 source_xe: &mut Source,
42 scratch: &mut Scratch<BE>,
43 ) where
44 S1: GLWESecretToRef,
45 S2: GLWESecretToRef,
46 M: GLWESwitchingKeyEncryptSk<BE>,
47 Scratch<BE>: ScratchTakeCore<BE>,
48 {
49 module.glwe_switching_key_encrypt_sk(self, sk_in, sk_out, source_xa, source_xe, scratch);
50 }
51}
52
53pub trait GLWESwitchingKeyEncryptSk<BE: Backend> {
54 fn glwe_switching_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
55 where
56 A: GGLWEInfos;
57
58 fn glwe_switching_key_encrypt_sk<R, S1, S2>(
59 &self,
60 res: &mut R,
61 sk_in: &S1,
62 sk_out: &S2,
63 source_xa: &mut Source,
64 source_xe: &mut Source,
65 scratch: &mut Scratch<BE>,
66 ) where
67 R: GGLWEToMut + GLWESwitchingKeyDegreesMut + GGLWEInfos,
68 S1: GLWESecretToRef,
69 S2: GLWESecretToRef;
70}
71
72impl<BE: Backend> GLWESwitchingKeyEncryptSk<BE> for Module<BE>
73where
74 Self: ModuleN + GGLWEEncryptSk<BE> + GLWESecretPreparedFactory<BE> + VecZnxSwitchRing + SvpPrepare<BE>,
75 Scratch<BE>: ScratchTakeCore<BE>,
76{
77 fn glwe_switching_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
78 where
79 A: GGLWEInfos,
80 {
81 self.gglwe_encrypt_sk_tmp_bytes(infos)
82 .max(ScalarZnx::bytes_of(self.n(), 1))
83 + ScalarZnx::bytes_of(self.n(), infos.rank_in().into())
84 + self.bytes_of_glwe_secret_prepared_from_infos(infos)
85 }
86
87 fn glwe_switching_key_encrypt_sk<R, S1, S2>(
88 &self,
89 res: &mut R,
90 sk_in: &S1,
91 sk_out: &S2,
92 source_xa: &mut Source,
93 source_xe: &mut Source,
94 scratch: &mut Scratch<BE>,
95 ) where
96 R: GGLWEToMut + GLWESwitchingKeyDegreesMut + GGLWEInfos,
97 S1: GLWESecretToRef,
98 S2: GLWESecretToRef,
99 {
100 let sk_in: &GLWESecret<&[u8]> = &sk_in.to_ref();
101 let sk_out: &GLWESecret<&[u8]> = &sk_out.to_ref();
102
103 assert!(sk_in.n().0 <= self.n() as u32);
104 assert!(sk_out.n().0 <= self.n() as u32);
105 assert!(
106 scratch.available() >= self.glwe_switching_key_encrypt_sk_tmp_bytes(res),
107 "scratch.available()={} < GLWESwitchingKey::encrypt_sk_tmp_bytes={}",
108 scratch.available(),
109 self.glwe_switching_key_encrypt_sk_tmp_bytes(res)
110 );
111
112 let (mut sk_in_tmp, scratch_1) = scratch.take_scalar_znx(self.n(), sk_in.rank().into());
113 for i in 0..sk_in.rank().into() {
114 self.vec_znx_switch_ring(
115 &mut sk_in_tmp.as_vec_znx_mut(),
116 i,
117 &sk_in.data.as_vec_znx(),
118 i,
119 );
120 }
121
122 let (mut sk_out_tmp, scratch_2) = scratch_1.take_glwe_secret_prepared(self, sk_out.rank());
123 {
124 let (mut tmp, _) = scratch_2.take_scalar_znx(self.n(), 1);
125 for i in 0..sk_out.rank().into() {
126 self.vec_znx_switch_ring(&mut tmp.as_vec_znx_mut(), 0, &sk_out.data.as_vec_znx(), i);
127 self.svp_prepare(&mut sk_out_tmp.data, i, &tmp, 0);
128 }
129 }
130
131 sk_out_tmp.dist = sk_out.dist;
132
133 self.gglwe_encrypt_sk(
134 res,
135 &sk_in_tmp,
136 &sk_out_tmp,
137 source_xa,
138 source_xe,
139 scratch_2,
140 );
141
142 *res.input_degree() = sk_in.n();
143 *res.output_degree() = sk_out.n();
144 }
145}
146
147pub trait GLWESwitchingKeyEncryptPk<BE: Backend> {
148 fn glwe_switching_key_encrypt_pk_tmp_bytes<A>(&self, infos: &A) -> usize
149 where
150 A: GGLWEInfos;
151}
152
153impl<BE: Backend> GLWESwitchingKeyEncryptPk<BE> for Module<BE> {
154 fn glwe_switching_key_encrypt_pk_tmp_bytes<A>(&self, _infos: &A) -> usize
155 where
156 A: GGLWEInfos,
157 {
158 unimplemented!()
159 }
160}