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
240impl<K, C, S> Encryptor<[u8], SeededBatchedFheUint8<C, S>> for K
241where
242 K: Encryptor<[bool], (Vec<C>, S)>,
243{
244 /// Encrypt a slice of u8s of arbitray length packed into collection of
245 /// seeded RLWE ciphertexts and return `SeededBatchedFheUint8`
246 fn encrypt(&self, m: &[u8]) -> SeededBatchedFheUint8<C, S> {
247 // convert vector of u8s to vector bools
248 let bool_m = m
249 .iter()
250 .flat_map(|v| (0..8).into_iter().map(|i| (((*v) >> i) & 1) == 1))
251 .collect_vec();
252 let (cts, seed) = K::encrypt(&self, &bool_m);
253 SeededBatchedFheUint8 {
254 data: cts,
255 seed,
256 count: m.len(),
257 }
258 }
259}
260
261impl<C, S> SeededBatchedFheUint8<C, S> {
262 /// Unseed collection of seeded RLWE ciphertexts of `Self` and returns
263 /// `NonInteractiveBatchedFheUint8` with collection of unseeded RLWE
264 /// ciphertexts.
265 ///
266 /// In non-interactive MPC setting, RLWE ciphertexts are encrypted under
267 /// user's RLWE secret `u_j`. The RLWE ciphertexts must be key switched to
268 /// ideal RLWE secret `s` of the MPC protocol before use.
269 ///
270 /// Note that we don't provide `unseed` API from `Self` to
271 /// `BatchedFheUint8`. This is because:
272 ///
273 /// - In non-interactive setting (1) client encrypts private inputs using
274 /// their secret `u_j` as `SeededBatchedFheUint8` and sends it to the
275 /// server. (2) Server unseeds `SeededBatchedFheUint8` into
276 /// `NonInteractiveBatchedFheUint8` indicating that private inputs are
277 /// still encrypted under user's RLWE secret `u_j`. (3) Server key
278 /// switches `NonInteractiveBatchedFheUint8` from user's RLWE secret `u_j`
279 /// to ideal RLWE secret `s` and outputs `BatchedFheUint8`. (4)
280 /// `BatchedFheUint8` always stores RLWE secret under ideal RLWE secret of
281 /// the protocol. Hence, it is safe to extract FheUint8s. Server proceeds
282 /// to extract necessary FheUint8s.
283 ///
284 /// - In interactive setting (1) client always encrypts private inputs using
285 /// public key corresponding to ideal RLWE secret `s` of the protocol and
286 /// produces `BatchedFheUint8`. (2) Given `BatchedFheUint8` stores
287 /// collection of RLWE ciphertext under ideal RLWE secret `s`, server can
288 /// directly extract necessary FheUint8s to use.
289 ///
290 /// Thus, there's no need to go directly from `Self` to `BatchedFheUint8`.
291 pub fn unseed<M>(&self) -> NonInteractiveBatchedFheUint8<M>
292 where
293 NonInteractiveBatchedFheUint8<M>: for<'a> From<&'a SeededBatchedFheUint8<C, S>>,
294 M: Matrix<R = C>,
295 {
296 NonInteractiveBatchedFheUint8::from(self)
297 }
298}
299
300impl<C, K> MultiPartyDecryptor<u8, FheUint8<C>> for K
301where
302 K: MultiPartyDecryptor<bool, C>,
303 <Self as MultiPartyDecryptor<bool, C>>::DecryptionShare: Clone,
304{
305 type DecryptionShare = Vec<<Self as MultiPartyDecryptor<bool, C>>::DecryptionShare>;
306 fn gen_decryption_share(&self, c: &FheUint8<C>) -> Self::DecryptionShare {
307 assert!(c.data().len() == 8);
308 c.data()
309 .iter()
310 .map(|bit_c| {
311 let decryption_share =
312 MultiPartyDecryptor::<bool, C>::gen_decryption_share(self, bit_c);
313 decryption_share
314 })
315 .collect_vec()
316 }
317
318 fn aggregate_decryption_shares(&self, c: &FheUint8<C>, shares: &[Self::DecryptionShare]) -> u8 {
319 let mut out = 0u8;
320
321 (0..8).into_iter().for_each(|i| {
322 // Collect bit i^th decryption share of each party
323 let bit_i_decryption_shares = shares.iter().map(|s| s[i].clone()).collect_vec();
324 let bit_i = MultiPartyDecryptor::<bool, C>::aggregate_decryption_shares(
325 self,
326 &c.data()[i],
327 &bit_i_decryption_shares,
328 );
329
330 if bit_i {
331 out += 1 << i;
332 }
333 });
334
335 out
336 }
337}
338
339impl<C, K> Encryptor<u8, FheUint8<C>> for K
340where
341 K: Encryptor<bool, C>,
342{
343 fn encrypt(&self, m: &u8) -> FheUint8<C> {
344 let cts = (0..8)
345 .into_iter()
346 .map(|i| {
347 let bit = ((m >> i) & 1) == 1;
348 K::encrypt(self, &bit)
349 })
350 .collect_vec();
351 FheUint8 { data: cts }
352 }
353}
354
355impl<K, C> Decryptor<u8, FheUint8<C>> for K
356where
357 K: Decryptor<bool, C>,
358{
359 fn decrypt(&self, c: &FheUint8<C>) -> u8 {
360 assert!(c.data.len() == 8);
361 let mut out = 0u8;
362 c.data().iter().enumerate().for_each(|(index, bit_c)| {
363 let bool = K::decrypt(self, bit_c);
364 if bool {
365 out += 1 << index;
366 }
367 });
368 out
369 }
370}