poulpy/core/src/glwe/public_key.rs
core/src/glwe/public_key.rsRaw
1use backend::hal::{
2 api::{
3 ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAlloc, VecZnxAllocBytes, VecZnxDftAlloc, VecZnxDftAllocBytes,
4 VecZnxDftFromVecZnx,
5 },
6 layouts::{Backend, Data, DataMut, DataRef, Module, ReaderFrom, Scratch, ScratchOwned, VecZnx, VecZnxDft, WriterTo},
7 oep::{ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeVecZnxDftImpl, TakeVecZnxImpl},
8};
9use sampling::source::Source;
10
11use crate::{GLWECiphertext, GLWEEncryptSkFamily, GLWESecretExec, Infos, dist::Distribution};
12
13pub trait GLWEPublicKeyFamily<B: Backend> = GLWEEncryptSkFamily<B>;
14
15#[derive(PartialEq, Eq)]
16pub struct GLWEPublicKey<D: Data> {
17 pub(crate) data: VecZnx<D>,
18 pub(crate) basek: usize,
19 pub(crate) k: usize,
20 pub(crate) dist: Distribution,
21}
22
23impl GLWEPublicKey<Vec<u8>> {
24 pub fn alloc<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank: usize) -> Self
25 where
26 Module<B>: VecZnxAlloc,
27 {
28 Self {
29 data: module.vec_znx_alloc(rank + 1, k.div_ceil(basek)),
30 basek: basek,
31 k: k,
32 dist: Distribution::NONE,
33 }
34 }
35
36 pub fn bytes_of<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize
37 where
38 Module<B>: VecZnxAllocBytes,
39 {
40 module.vec_znx_alloc_bytes(rank + 1, k.div_ceil(basek))
41 }
42}
43
44impl<D: Data> Infos for GLWEPublicKey<D> {
45 type Inner = VecZnx<D>;
46
47 fn inner(&self) -> &Self::Inner {
48 &self.data
49 }
50
51 fn basek(&self) -> usize {
52 self.basek
53 }
54
55 fn k(&self) -> usize {
56 self.k
57 }
58}
59
60impl<D: Data> GLWEPublicKey<D> {
61 pub fn rank(&self) -> usize {
62 self.cols() - 1
63 }
64}
65
66impl<D: DataMut> GLWEPublicKey<D> {
67 pub fn generate_from_sk<S: DataRef, B: Backend>(
68 &mut self,
69 module: &Module<B>,
70 sk: &GLWESecretExec<S, B>,
71 source_xa: &mut Source,
72 source_xe: &mut Source,
73 sigma: f64,
74 ) where
75 Module<B>: GLWEPublicKeyFamily<B> + VecZnxAlloc,
76 B: ScratchOwnedAllocImpl<B>
77 + ScratchOwnedBorrowImpl<B>
78 + TakeVecZnxDftImpl<B>
79 + ScratchAvailableImpl<B>
80 + TakeVecZnxImpl<B>,
81 {
82 #[cfg(debug_assertions)]
83 {
84 match sk.dist {
85 Distribution::NONE => panic!("invalid sk: SecretDistribution::NONE"),
86 _ => {}
87 }
88 }
89
90 // Its ok to allocate scratch space here since pk is usually generated only once.
91 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GLWECiphertext::encrypt_sk_scratch_space(
92 module,
93 self.basek(),
94 self.k(),
95 ));
96
97 let mut tmp: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(module, self.basek(), self.k(), self.rank());
98 tmp.encrypt_zero_sk(module, sk, source_xa, source_xe, sigma, scratch.borrow());
99 self.dist = sk.dist;
100 }
101}
102
103use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
104
105impl<D: DataMut> ReaderFrom for GLWEPublicKey<D> {
106 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
107 self.k = reader.read_u64::<LittleEndian>()? as usize;
108 self.basek = reader.read_u64::<LittleEndian>()? as usize;
109 match Distribution::read_from(reader) {
110 Ok(dist) => self.dist = dist,
111 Err(e) => return Err(e),
112 }
113 self.data.read_from(reader)
114 }
115}
116
117impl<D: DataRef> WriterTo for GLWEPublicKey<D> {
118 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
119 writer.write_u64::<LittleEndian>(self.k as u64)?;
120 writer.write_u64::<LittleEndian>(self.basek as u64)?;
121 match self.dist.write_to(writer) {
122 Ok(()) => {}
123 Err(e) => return Err(e),
124 }
125 self.data.write_to(writer)
126 }
127}
128
129#[derive(PartialEq, Eq)]
130pub struct GLWEPublicKeyExec<D: Data, B: Backend> {
131 pub(crate) data: VecZnxDft<D, B>,
132 pub(crate) basek: usize,
133 pub(crate) k: usize,
134 pub(crate) dist: Distribution,
135}
136
137impl<D: Data, B: Backend> Infos for GLWEPublicKeyExec<D, B> {
138 type Inner = VecZnxDft<D, B>;
139
140 fn inner(&self) -> &Self::Inner {
141 &self.data
142 }
143
144 fn basek(&self) -> usize {
145 self.basek
146 }
147
148 fn k(&self) -> usize {
149 self.k
150 }
151}
152
153impl<D: Data, B: Backend> GLWEPublicKeyExec<D, B> {
154 pub fn rank(&self) -> usize {
155 self.cols() - 1
156 }
157}
158
159impl<B: Backend> GLWEPublicKeyExec<Vec<u8>, B> {
160 pub fn alloc(module: &Module<B>, basek: usize, k: usize, rank: usize) -> Self
161 where
162 Module<B>: VecZnxDftAlloc<B>,
163 {
164 Self {
165 data: module.vec_znx_dft_alloc(rank + 1, k.div_ceil(basek)),
166 basek: basek,
167 k: k,
168 dist: Distribution::NONE,
169 }
170 }
171
172 pub fn bytes_of(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize
173 where
174 Module<B>: VecZnxDftAllocBytes,
175 {
176 module.vec_znx_dft_alloc_bytes(rank + 1, k.div_ceil(basek))
177 }
178
179 pub fn from<DataOther>(module: &Module<B>, other: &GLWEPublicKey<DataOther>, scratch: &mut Scratch<B>) -> Self
180 where
181 DataOther: DataRef,
182 Module<B>: VecZnxDftAlloc<B> + VecZnxDftFromVecZnx<B>,
183 {
184 let mut pk_exec: GLWEPublicKeyExec<Vec<u8>, B> = GLWEPublicKeyExec::alloc(module, other.basek(), other.k(), other.rank());
185 pk_exec.prepare(module, other, scratch);
186 pk_exec
187 }
188}
189
190impl<D: DataMut, B: Backend> GLWEPublicKeyExec<D, B> {
191 pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &GLWEPublicKey<DataOther>, _scratch: &mut Scratch<B>)
192 where
193 DataOther: DataRef,
194 Module<B>: VecZnxDftFromVecZnx<B>,
195 {
196 #[cfg(debug_assertions)]
197 {
198 assert_eq!(self.n(), module.n());
199 assert_eq!(other.n(), module.n());
200 assert_eq!(self.size(), other.size());
201 }
202
203 (0..self.cols()).for_each(|i| {
204 module.vec_znx_dft_from_vec_znx(1, 0, &mut self.data, i, &other.data, i);
205 });
206 self.k = other.k;
207 self.basek = other.basek;
208 self.dist = other.dist;
209 }
210}