phantom-zone/src/shortint/enc_dec.rs
src/shortint/enc_dec.rsRaw
1use itertools::Itertools;
2
3use crate::{
4 bool::BoolEvaluator,
5 random::{DefaultSecureRng, RandomFillUniformInModulus},
6 utils::WithLocal,
7 Decryptor, Encryptor, KeySwitchWithId, Matrix, MatrixEntity, MatrixMut, MultiPartyDecryptor,
8 RowMut, SampleExtractor,
9};
10
11/// Fhe UInt8
12///
13/// Note that `Self.data` stores encryptions of bits in little endian (i.e least
14/// signficant bit stored at 0th index and most signficant bit stores at 7th
15/// index)
16#[derive(Clone)]
17pub struct FheUint8<C> {
18 pub(super) data: Vec<C>,
19}
20
21impl<C> FheUint8<C> {
22 pub(super) fn data(&self) -> &[C] {
23 &self.data
24 }
25
26 pub(super) fn data_mut(&mut self) -> &mut [C] {
27 &mut self.data
28 }
29}
30
31/// Stores a batch of Fhe Uint8 ciphertext as collection of unseeded RLWE
32/// ciphertexts always encrypted under the ideal RLWE secret `s` of the MPC
33/// protocol
34///
35/// To extract Fhe Uint8 ciphertext at `index` call `self.extract(index)`
36pub struct BatchedFheUint8<C> {
37 /// Vector of RLWE ciphertexts `C`
38 data: Vec<C>,
39 /// Count of FheUint8s packed in vector of RLWE ciphertexts
40 count: usize,
41}
42
43impl<K, C> Encryptor<[u8], BatchedFheUint8<C>> for K
44where
45 K: Encryptor<[bool], Vec<C>>,
46{
47 /// Encrypt a batch of uint8s packed in vector of RLWE ciphertexts
48 ///
49 /// Uint8s can be extracted from `BatchedFheUint8` with `SampleExtractor`
50 fn encrypt(&self, m: &[u8]) -> BatchedFheUint8<C> {
51 let bool_m = m
52 .iter()
53 .flat_map(|v| {
54 (0..8)
55 .into_iter()
56 .map(|i| ((*v >> i) & 1) == 1)
57 .collect_vec()
58 })
59 .collect_vec();
60 let cts = K::encrypt(&self, &bool_m);
61 BatchedFheUint8 {
62 data: cts,
63 count: m.len(),
64 }
65 }
66}
67
68impl<M: MatrixEntity + MatrixMut<MatElement = u64>> From<&SeededBatchedFheUint8<M::R, [u8; 32]>>
69 for BatchedFheUint8<M>
70where
71 <M as Matrix>::R: RowMut,
72{
73 /// Unseeds collection of seeded RLWE ciphertext in `SeededBatchedFheUint8`
74 /// and returns as `Self`
75 fn from(value: &SeededBatchedFheUint8<M::R, [u8; 32]>) -> Self {
76 BoolEvaluator::with_local(|e| {
77 let parameters = e.parameters();
78 let ring_size = parameters.rlwe_n().0;
79 let rlwe_q = parameters.rlwe_q();
80
81 let mut prng = DefaultSecureRng::new_seeded(value.seed);
82 let rlwes = value
83 .data
84 .iter()
85 .map(|partb| {
86 let mut rlwe = M::zeros(2, ring_size);
87
88 // sample A
89 RandomFillUniformInModulus::random_fill(&mut prng, rlwe_q, rlwe.get_row_mut(0));
90
91 // Copy over B
92 rlwe.get_row_mut(1).copy_from_slice(partb.as_ref());
93
94 rlwe
95 })
96 .collect_vec();
97 Self {
98 data: rlwes,
99 count: value.count,
100 }
101 })
102 }
103}
104
105impl<C, R> SampleExtractor<FheUint8<R>> for BatchedFheUint8<C>
106where
107 C: SampleExtractor<R>,
108{
109 /// Extract Fhe Uint8 ciphertext at `index`
110 ///
111 /// `Self` stores batch of Fhe uint8 ciphertext as vector of RLWE
112 /// ciphertexts. Since Fhe uint8 ciphertext is collection of 8 bool
113 /// ciphertexts, Fhe uint8 ciphertext at index `i` is stored in coefficients
114 /// `i*8...(i+1)*8`. To extract Fhe uint8 at index `i`, sample extract bool
115 /// ciphertext at indices `[i*8, ..., (i+1)*8)`
116 fn extract_at(&self, index: usize) -> FheUint8<R> {
117 assert!(index < self.count);
118 BoolEvaluator::with_local(|e| {
119 let ring_size = e.parameters().rlwe_n().0;
120
121 let start_index = index * 8;
122 let end_index = (index + 1) * 8;
123 let data = (start_index..end_index)
124 .map(|i| {
125 let rlwe_index = i / ring_size;
126 let coeff_index = i % ring_size;
127 self.data[rlwe_index].extract_at(coeff_index)
128 })
129 .collect_vec();
130 FheUint8 { data }
131 })
132 }
133
134 /// Extracts all FheUint8s packed in vector of RLWE ciphertexts of `Self`
135 fn extract_all(&self) -> Vec<FheUint8<R>> {
136 (0..self.count)
137 .map(|index| self.extract_at(index))
138 .collect_vec()
139 }
140
141 /// Extracts first `how_many` FheUint8s packed in vector of RLWE
142 /// ciphertexts of `Self`
143 fn extract_many(&self, how_many: usize) -> Vec<FheUint8<R>> {
144 (0..how_many)
145 .map(|index| self.extract_at(index))
146 .collect_vec()
147 }
148}
149
150/// Stores a batch of FheUint8s packed in a collection unseeded RLWE ciphertexts
151///
152/// `Self` stores unseeded RLWE ciphertexts encrypted under user's RLWE secret
153/// `u_j` and is different from `BatchFheUint8` which stores collection of RLWE
154/// ciphertexts under ideal RLWE secret `s` of the (non-interactive/interactive)
155/// MPC protocol.
156///
157/// To extract FheUint8s from `Self`'s collection of RLWE ciphertexts, first
158/// switch `Self` to `BatchFheUint8` with `key_switch(user_id)` where `user_id`
159/// is user's id. This key switches collection of RLWE ciphertexts from
160/// user's RLWE secret `u_j` to ideal RLWE secret `s` of the MPC protocol. Then
161/// proceed to use `SampleExtract` on `BatchFheUint8` (for ex, call
162/// `extract_at(0)` to extract FheUint8 stored at index 0)
163pub struct NonInteractiveBatchedFheUint8<C> {
164 /// Vector of RLWE ciphertexts `C`
165 data: Vec<C>,
166 /// Count of FheUint8s packed in vector of RLWE ciphertexts
167 count: usize,
168}
169
170impl<M: MatrixEntity + MatrixMut<MatElement = u64>> From<&SeededBatchedFheUint8<M::R, [u8; 32]>>
171 for NonInteractiveBatchedFheUint8<M>
172where
173 <M as Matrix>::R: RowMut,
174{
175 /// Unseeds collection of seeded RLWE ciphertext in SeededBatchedFheUint8
176 /// and returns as `Self`
177 fn from(value: &SeededBatchedFheUint8<M::R, [u8; 32]>) -> Self {
178 BoolEvaluator::with_local(|e| {
179 let parameters = e.parameters();
180 let ring_size = parameters.rlwe_n().0;
181 let rlwe_q = parameters.rlwe_q();
182
183 let mut prng = DefaultSecureRng::new_seeded(value.seed);
184 let rlwes = value
185 .data
186 .iter()
187 .map(|partb| {
188 let mut rlwe = M::zeros(2, ring_size);
189
190 // sample A
191 RandomFillUniformInModulus::random_fill(&mut prng, rlwe_q, rlwe.get_row_mut(0));
192
193 // Copy over B
194 rlwe.get_row_mut(1).copy_from_slice(partb.as_ref());
195
196 rlwe
197 })
198 .collect_vec();
199 Self {
200 data: rlwes,
201 count: value.count,
202 }
203 })
204 }
205}
206
207impl<C> KeySwitchWithId<BatchedFheUint8<C>> for NonInteractiveBatchedFheUint8<C>
208where
209 C: KeySwitchWithId<C>,
210{
211 /// Key switch `Self`'s collection of RLWE cihertexts encrypted under user's
212 /// RLWE secret `u_j` to ideal RLWE secret `s` of the MPC protocol.
213 ///
214 /// - user_id: user id of user `j`
215 fn key_switch(&self, user_id: usize) -> BatchedFheUint8<C> {
216 let data = self
217 .data
218 .iter()
219 .map(|c| c.key_switch(user_id))
220 .collect_vec();
221 BatchedFheUint8 {
222 data,
223 count: self.count,
224 }
225 }
226}
227
228pub struct SeededBatchedFheUint8<C, S> {
229 /// Vector of Seeded RLWE ciphertexts `C`.
230 ///
231 /// If RLWE(m) = [a, b] s.t. m + e = b - as, `a` can be seeded and seeded
232 /// RLWE ciphertext only contains `b` polynomial
233 data: Vec<C>,
234 /// Seed for the ciphertexts
235 seed: S,
236 /// Count of FheUint8s packed in vector of RLWE ciphertexts
237 count: usize,
238}
239
240#[cfg(feature = "non_interactive_mp")]
241impl<K, C, S> Encryptor<[u8], SeededBatchedFheUint8<C, S>> for K
242where
243 K: Encryptor<[bool], crate::NonInteractiveSeededFheBools<C, S>>,
244{
245 /// Encrypt a slice of u8s of arbitray length packed into collection of
246 /// seeded RLWE ciphertexts and return `SeededBatchedFheUint8`
247 fn encrypt(&self, m: &[u8]) -> SeededBatchedFheUint8<C, S> {
248 // convert vector of u8s to vector bools
249 let bool_m = m
250 .iter()
251 .flat_map(|v| (0..8).into_iter().map(|i| (((*v) >> i) & 1) == 1))
252 .collect_vec();
253 let (cts, seed): (Vec<C>, S) = K::encrypt(&self, &bool_m).into();
254 SeededBatchedFheUint8 {
255 data: cts,
256 seed,
257 count: m.len(),
258 }
259 }
260}
261
262impl<C, S> SeededBatchedFheUint8<C, S> {
263 /// Unseed collection of seeded RLWE ciphertexts of `Self` and returns
264 /// `NonInteractiveBatchedFheUint8` with collection of unseeded RLWE
265 /// ciphertexts.
266 ///
267 /// In non-interactive MPC setting, RLWE ciphertexts are encrypted under
268 /// user's RLWE secret `u_j`. The RLWE ciphertexts must be key switched to
269 /// ideal RLWE secret `s` of the MPC protocol before use.
270 ///
271 /// Note that we don't provide `unseed` API from `Self` to
272 /// `BatchedFheUint8`. This is because:
273 ///
274 /// - In non-interactive setting (1) client encrypts private inputs using
275 /// their secret `u_j` as `SeededBatchedFheUint8` and sends it to the
276 /// server. (2) Server unseeds `SeededBatchedFheUint8` into
277 /// `NonInteractiveBatchedFheUint8` indicating that private inputs are
278 /// still encrypted under user's RLWE secret `u_j`. (3) Server key
279 /// switches `NonInteractiveBatchedFheUint8` from user's RLWE secret `u_j`
280 /// to ideal RLWE secret `s` and outputs `BatchedFheUint8`. (4)
281 /// `BatchedFheUint8` always stores RLWE secret under ideal RLWE secret of
282 /// the protocol. Hence, it is safe to extract FheUint8s. Server proceeds
283 /// to extract necessary FheUint8s.
284 ///
285 /// - In interactive setting (1) client always encrypts private inputs using
286 /// public key corresponding to ideal RLWE secret `s` of the protocol and
287 /// produces `BatchedFheUint8`. (2) Given `BatchedFheUint8` stores
288 /// collection of RLWE ciphertext under ideal RLWE secret `s`, server can
289 /// directly extract necessary FheUint8s to use.
290 ///
291 /// Thus, there's no need to go directly from `Self` to `BatchedFheUint8`.
292 pub fn unseed<M>(&self) -> NonInteractiveBatchedFheUint8<M>
293 where
294 NonInteractiveBatchedFheUint8<M>: for<'a> From<&'a SeededBatchedFheUint8<C, S>>,
295 M: Matrix<R = C>,
296 {
297 NonInteractiveBatchedFheUint8::from(self)
298 }
299}
300
301impl<C, K> MultiPartyDecryptor<u8, FheUint8<C>> for K
302where
303 K: MultiPartyDecryptor<bool, C>,
304 <Self as MultiPartyDecryptor<bool, C>>::DecryptionShare: Clone,
305{
306 type DecryptionShare = Vec<<Self as MultiPartyDecryptor<bool, C>>::DecryptionShare>;
307 fn gen_decryption_share(&self, c: &FheUint8<C>) -> Self::DecryptionShare {
308 assert!(c.data().len() == 8);
309 c.data()
310 .iter()
311 .map(|bit_c| {
312 let decryption_share =
313 MultiPartyDecryptor::<bool, C>::gen_decryption_share(self, bit_c);
314 decryption_share
315 })
316 .collect_vec()
317 }
318
319 fn aggregate_decryption_shares(&self, c: &FheUint8<C>, shares: &[Self::DecryptionShare]) -> u8 {
320 let mut out = 0u8;
321
322 (0..8).into_iter().for_each(|i| {
323 // Collect bit i^th decryption share of each party
324 let bit_i_decryption_shares = shares.iter().map(|s| s[i].clone()).collect_vec();
325 let bit_i = MultiPartyDecryptor::<bool, C>::aggregate_decryption_shares(
326 self,
327 &c.data()[i],
328 &bit_i_decryption_shares,
329 );
330
331 if bit_i {
332 out += 1 << i;
333 }
334 });
335
336 out
337 }
338}
339
340impl<C, K> Encryptor<u8, FheUint8<C>> for K
341where
342 K: Encryptor<bool, C>,
343{
344 fn encrypt(&self, m: &u8) -> FheUint8<C> {
345 let cts = (0..8)
346 .into_iter()
347 .map(|i| {
348 let bit = ((m >> i) & 1) == 1;
349 K::encrypt(self, &bit)
350 })
351 .collect_vec();
352 FheUint8 { data: cts }
353 }
354}
355
356impl<K, C> Decryptor<u8, FheUint8<C>> for K
357where
358 K: Decryptor<bool, C>,
359{
360 fn decrypt(&self, c: &FheUint8<C>) -> u8 {
361 assert!(c.data.len() == 8);
362 let mut out = 0u8;
363 c.data().iter().enumerate().for_each(|(index, bit_c)| {
364 let bool = K::decrypt(self, bit_c);
365 if bool {
366 out += 1 << index;
367 }
368 });
369 out
370 }
371}