poulpy/core/src/keys.rs
core/src/keys.rsRaw
1use base2k::{
2 Backend, FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScalarZnxDft, ScalarZnxDftAlloc, ScalarZnxDftOps, ScalarZnxDftToMut,
3 ScalarZnxDftToRef, ScalarZnxToMut, ScalarZnxToRef, ScratchOwned, VecZnxDft, VecZnxDftToMut, VecZnxDftToRef, ZnxInfos,
4 ZnxZero,
5};
6use sampling::source::Source;
7
8use crate::{elem::Infos, glwe_ciphertext_fourier::GLWECiphertextFourier};
9
10#[derive(Clone, Copy, Debug)]
11pub enum SecretDistribution {
12 TernaryFixed(usize), // Ternary with fixed Hamming weight
13 TernaryProb(f64), // Ternary with probabilistic Hamming weight
14 ZERO, // Debug mod
15 NONE,
16}
17
18pub struct SecretKey<T> {
19 pub data: ScalarZnx<T>,
20 pub dist: SecretDistribution,
21}
22
23impl SecretKey<Vec<u8>> {
24 pub fn new<B: Backend>(module: &Module<B>, rank: usize) -> Self {
25 Self {
26 data: module.new_scalar_znx(rank),
27 dist: SecretDistribution::NONE,
28 }
29 }
30}
31
32impl<DataSelf> SecretKey<DataSelf> {
33 pub fn n(&self) -> usize {
34 self.data.n()
35 }
36
37 pub fn log_n(&self) -> usize {
38 self.data.log_n()
39 }
40
41 pub fn rank(&self) -> usize {
42 self.data.cols()
43 }
44}
45
46impl<S> SecretKey<S>
47where
48 S: AsMut<[u8]> + AsRef<[u8]>,
49{
50 pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
51 (0..self.rank()).for_each(|i| {
52 self.data.fill_ternary_prob(i, prob, source);
53 });
54 self.dist = SecretDistribution::TernaryProb(prob);
55 }
56
57 pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
58 (0..self.rank()).for_each(|i| {
59 self.data.fill_ternary_hw(i, hw, source);
60 });
61 self.dist = SecretDistribution::TernaryFixed(hw);
62 }
63
64 pub fn fill_zero(&mut self) {
65 self.data.zero();
66 self.dist = SecretDistribution::ZERO;
67 }
68}
69
70impl<C> ScalarZnxToMut for SecretKey<C>
71where
72 ScalarZnx<C>: ScalarZnxToMut,
73{
74 fn to_mut(&mut self) -> ScalarZnx<&mut [u8]> {
75 self.data.to_mut()
76 }
77}
78
79impl<C> ScalarZnxToRef for SecretKey<C>
80where
81 ScalarZnx<C>: ScalarZnxToRef,
82{
83 fn to_ref(&self) -> ScalarZnx<&[u8]> {
84 self.data.to_ref()
85 }
86}
87
88pub struct SecretKeyFourier<T, B: Backend> {
89 pub data: ScalarZnxDft<T, B>,
90 pub dist: SecretDistribution,
91}
92
93impl<DataSelf, B: Backend> SecretKeyFourier<DataSelf, B> {
94 pub fn n(&self) -> usize {
95 self.data.n()
96 }
97
98 pub fn log_n(&self) -> usize {
99 self.data.log_n()
100 }
101
102 pub fn rank(&self) -> usize {
103 self.data.cols()
104 }
105}
106
107impl<B: Backend> SecretKeyFourier<Vec<u8>, B> {
108 pub fn new(module: &Module<B>, rank: usize) -> Self {
109 Self {
110 data: module.new_scalar_znx_dft(rank),
111 dist: SecretDistribution::NONE,
112 }
113 }
114
115 pub fn dft<S>(&mut self, module: &Module<FFT64>, sk: &SecretKey<S>)
116 where
117 SecretKeyFourier<Vec<u8>, B>: ScalarZnxDftToMut<base2k::FFT64>,
118 SecretKey<S>: ScalarZnxToRef,
119 {
120 #[cfg(debug_assertions)]
121 {
122 match sk.dist {
123 SecretDistribution::NONE => panic!("invalid sk: SecretDistribution::NONE"),
124 _ => {}
125 }
126
127 assert_eq!(self.n(), module.n());
128 assert_eq!(sk.n(), module.n());
129 assert_eq!(self.rank(), sk.rank());
130 }
131
132 (0..self.rank()).for_each(|i| {
133 module.svp_prepare(self, i, sk, i);
134 });
135 self.dist = sk.dist;
136 }
137}
138
139impl<C, B: Backend> ScalarZnxDftToMut<B> for SecretKeyFourier<C, B>
140where
141 ScalarZnxDft<C, B>: ScalarZnxDftToMut<B>,
142{
143 fn to_mut(&mut self) -> ScalarZnxDft<&mut [u8], B> {
144 self.data.to_mut()
145 }
146}
147
148impl<C, B: Backend> ScalarZnxDftToRef<B> for SecretKeyFourier<C, B>
149where
150 ScalarZnxDft<C, B>: ScalarZnxDftToRef<B>,
151{
152 fn to_ref(&self) -> ScalarZnxDft<&[u8], B> {
153 self.data.to_ref()
154 }
155}
156
157pub struct GLWEPublicKey<D, B: Backend> {
158 pub data: GLWECiphertextFourier<D, B>,
159 pub dist: SecretDistribution,
160}
161
162impl<B: Backend> GLWEPublicKey<Vec<u8>, B> {
163 pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize, rank: usize) -> Self {
164 Self {
165 data: GLWECiphertextFourier::new(module, log_base2k, log_k, rank),
166 dist: SecretDistribution::NONE,
167 }
168 }
169}
170
171impl<T, B: Backend> Infos for GLWEPublicKey<T, B> {
172 type Inner = VecZnxDft<T, B>;
173
174 fn inner(&self) -> &Self::Inner {
175 &self.data.data
176 }
177
178 fn basek(&self) -> usize {
179 self.data.basek
180 }
181
182 fn k(&self) -> usize {
183 self.data.k
184 }
185}
186
187impl<T, B: Backend> GLWEPublicKey<T, B> {
188 pub fn rank(&self) -> usize {
189 self.cols() - 1
190 }
191}
192
193impl<C, B: Backend> VecZnxDftToMut<B> for GLWEPublicKey<C, B>
194where
195 VecZnxDft<C, B>: VecZnxDftToMut<B>,
196{
197 fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
198 self.data.to_mut()
199 }
200}
201
202impl<C, B: Backend> VecZnxDftToRef<B> for GLWEPublicKey<C, B>
203where
204 VecZnxDft<C, B>: VecZnxDftToRef<B>,
205{
206 fn to_ref(&self) -> VecZnxDft<&[u8], B> {
207 self.data.to_ref()
208 }
209}
210
211impl<C> GLWEPublicKey<C, FFT64> {
212 pub fn generate<S>(
213 &mut self,
214 module: &Module<FFT64>,
215 sk_dft: &SecretKeyFourier<S, FFT64>,
216 source_xa: &mut Source,
217 source_xe: &mut Source,
218 sigma: f64,
219 bound: f64,
220 ) where
221 VecZnxDft<C, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64>,
222 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64> + ZnxInfos,
223 {
224 #[cfg(debug_assertions)]
225 {
226 match sk_dft.dist {
227 SecretDistribution::NONE => panic!("invalid sk_dft: SecretDistribution::NONE"),
228 _ => {}
229 }
230 }
231
232 // Its ok to allocate scratch space here since pk is usually generated only once.
233 let mut scratch: ScratchOwned = ScratchOwned::new(GLWECiphertextFourier::encrypt_sk_scratch_space(
234 module,
235 self.rank(),
236 self.size(),
237 ));
238 self.data.encrypt_zero_sk(
239 module,
240 sk_dft,
241 source_xa,
242 source_xe,
243 sigma,
244 bound,
245 scratch.borrow(),
246 );
247 self.dist = sk_dft.dist;
248 }
249}