phantom-zone/src/shortint/mod.rs
src/shortint/mod.rsRaw
1use itertools::Itertools;
2
3use crate::{
4 bool::evaluator::{
5 BoolEvaluator, ClientKey, PublicKey, ServerKeyEvaluationDomain, BOOL_SERVER_KEY,
6 },
7 utils::{Global, WithLocal},
8 Decryptor, Encryptor, Matrix, MultiPartyDecryptor,
9};
10
11mod ops;
12mod types;
13
14type FheUint8 = types::FheUint8<Vec<u64>>;
15
16impl Encryptor<u8, FheUint8> for ClientKey {
17 fn encrypt(&self, m: &u8) -> FheUint8 {
18 let cts = (0..8)
19 .into_iter()
20 .map(|i| {
21 let bit = ((m >> i) & 1) == 1;
22 Encryptor::<bool, Vec<u64>>::encrypt(self, &bit)
23 })
24 .collect_vec();
25 FheUint8 { data: cts }
26 }
27}
28
29impl Decryptor<u8, FheUint8> for ClientKey {
30 fn decrypt(&self, c: &FheUint8) -> u8 {
31 assert!(c.data.len() == 8);
32 let mut out = 0u8;
33 c.data().iter().enumerate().for_each(|(index, bit_c)| {
34 let bool = Decryptor::<bool, Vec<u64>>::decrypt(self, bit_c);
35 if bool {
36 out += 1 << index;
37 }
38 });
39 out
40 }
41}
42
43impl<M, R, Mo> Encryptor<u8, FheUint8> for PublicKey<M, R, Mo>
44where
45 PublicKey<M, R, Mo>: Encryptor<bool, Vec<u64>>,
46{
47 fn encrypt(&self, m: &u8) -> FheUint8 {
48 let cts = (0..8)
49 .into_iter()
50 .map(|i| {
51 let bit = ((m >> i) & 1) == 1;
52 Encryptor::<bool, Vec<u64>>::encrypt(self, &bit)
53 })
54 .collect_vec();
55 FheUint8 { data: cts }
56 }
57}
58
59impl MultiPartyDecryptor<u8, FheUint8> for ClientKey
60where
61 ClientKey: MultiPartyDecryptor<bool, Vec<u64>>,
62{
63 type DecryptionShare = Vec<<Self as MultiPartyDecryptor<bool, Vec<u64>>>::DecryptionShare>;
64 fn gen_decryption_share(&self, c: &FheUint8) -> Self::DecryptionShare {
65 assert!(c.data().len() == 8);
66 c.data()
67 .iter()
68 .map(|bit_c| {
69 let decryption_share =
70 MultiPartyDecryptor::<bool, Vec<u64>>::gen_decryption_share(self, bit_c);
71 decryption_share
72 })
73 .collect_vec()
74 }
75
76 fn aggregate_decryption_shares(&self, c: &FheUint8, shares: &[Self::DecryptionShare]) -> u8 {
77 let mut out = 0u8;
78
79 (0..8).into_iter().for_each(|i| {
80 // Collect bit i^th decryption share of each party
81 let bit_i_decryption_shares = shares.iter().map(|s| s[i]).collect_vec();
82 let bit_i = MultiPartyDecryptor::<bool, Vec<u64>>::aggregate_decryption_shares(
83 self,
84 &c.data()[i],
85 &bit_i_decryption_shares,
86 );
87
88 if bit_i {
89 out += 1 << i;
90 }
91 });
92
93 out
94 }
95}
96
97mod frontend {
98 use super::ops::{
99 arbitrary_bit_adder, arbitrary_bit_division_for_quotient_and_rem, arbitrary_bit_subtractor,
100 eight_bit_mul,
101 };
102 use crate::{
103 bool::evaluator::{BoolEvaluator, ServerKeyEvaluationDomain},
104 utils::{Global, WithLocal},
105 };
106
107 use super::FheUint8;
108
109 mod arithetic {
110 use crate::bool::{evaluator::BooleanGates, FheBool};
111
112 use super::*;
113 use std::ops::{Add, AddAssign, Div, Mul, Rem, Sub};
114
115 impl AddAssign<&FheUint8> for FheUint8 {
116 fn add_assign(&mut self, rhs: &FheUint8) {
117 BoolEvaluator::with_local_mut_mut(&mut |e| {
118 let key = ServerKeyEvaluationDomain::global();
119 arbitrary_bit_adder(e, self.data_mut(), rhs.data(), false, key);
120 });
121 }
122 }
123
124 impl Add<&FheUint8> for &FheUint8 {
125 type Output = FheUint8;
126 fn add(self, rhs: &FheUint8) -> Self::Output {
127 let mut a = self.clone();
128 a += rhs;
129 a
130 }
131 }
132
133 impl Sub<&FheUint8> for &FheUint8 {
134 type Output = FheUint8;
135 fn sub(self, rhs: &FheUint8) -> Self::Output {
136 BoolEvaluator::with_local_mut(|e| {
137 let key = ServerKeyEvaluationDomain::global();
138 let (out, _, _) = arbitrary_bit_subtractor(e, self.data(), rhs.data(), key);
139 FheUint8 { data: out }
140 })
141 }
142 }
143
144 impl Mul<&FheUint8> for &FheUint8 {
145 type Output = FheUint8;
146 fn mul(self, rhs: &FheUint8) -> Self::Output {
147 BoolEvaluator::with_local_mut(|e| {
148 let key = ServerKeyEvaluationDomain::global();
149 let out = eight_bit_mul(e, self.data(), rhs.data(), key);
150 FheUint8 { data: out }
151 })
152 }
153 }
154
155 impl Div<&FheUint8> for &FheUint8 {
156 type Output = FheUint8;
157 fn div(self, rhs: &FheUint8) -> Self::Output {
158 // TODO(Jay:) Figure out how to set zero error flag
159 BoolEvaluator::with_local_mut(|e| {
160 let key = ServerKeyEvaluationDomain::global();
161 let (quotient, _) = arbitrary_bit_division_for_quotient_and_rem(
162 e,
163 self.data(),
164 rhs.data(),
165 key,
166 );
167 FheUint8 { data: quotient }
168 })
169 }
170 }
171
172 impl Rem<&FheUint8> for &FheUint8 {
173 type Output = FheUint8;
174 fn rem(self, rhs: &FheUint8) -> Self::Output {
175 BoolEvaluator::with_local_mut(|e| {
176 let key = ServerKeyEvaluationDomain::global();
177 let (_, remainder) = arbitrary_bit_division_for_quotient_and_rem(
178 e,
179 self.data(),
180 rhs.data(),
181 key,
182 );
183 FheUint8 { data: remainder }
184 })
185 }
186 }
187
188 impl FheUint8 {
189 pub fn overflowing_add_assign(&mut self, rhs: &FheUint8) -> FheBool {
190 BoolEvaluator::with_local_mut_mut(&mut |e| {
191 let key = ServerKeyEvaluationDomain::global();
192 let (overflow, _) =
193 arbitrary_bit_adder(e, self.data_mut(), rhs.data(), false, key);
194 overflow
195 })
196 }
197
198 pub fn overflowing_add(self, rhs: &FheUint8) -> (FheUint8, FheBool) {
199 BoolEvaluator::with_local_mut(|e| {
200 let mut lhs = self.clone();
201 let key = ServerKeyEvaluationDomain::global();
202 let (overflow, _) =
203 arbitrary_bit_adder(e, lhs.data_mut(), rhs.data(), false, key);
204 (lhs, overflow)
205 })
206 }
207
208 pub fn overflowing_sub(&self, rhs: &FheUint8) -> (FheUint8, FheBool) {
209 BoolEvaluator::with_local_mut(|e| {
210 let key = ServerKeyEvaluationDomain::global();
211 let (out, mut overflow, _) =
212 arbitrary_bit_subtractor(e, self.data(), rhs.data(), key);
213 e.not_inplace(&mut overflow);
214 (FheUint8 { data: out }, overflow)
215 })
216 }
217
218 pub fn div_rem(&self, rhs: &FheUint8) -> (FheUint8, FheUint8) {
219 // TODO(Jay:) Figure out how to set zero error flag
220 BoolEvaluator::with_local_mut(|e| {
221 let key = ServerKeyEvaluationDomain::global();
222 let (quotient, remainder) = arbitrary_bit_division_for_quotient_and_rem(
223 e,
224 self.data(),
225 rhs.data(),
226 key,
227 );
228 (FheUint8 { data: quotient }, FheUint8 { data: remainder })
229 })
230 }
231 }
232 }
233
234 mod booleans {
235 use crate::{
236 bool::{evaluator::BooleanGates, FheBool},
237 shortint::ops::{
238 arbitrary_bit_comparator, arbitrary_bit_equality, arbitrary_signed_bit_comparator,
239 },
240 };
241
242 use super::*;
243
244 impl FheUint8 {
245 /// a == b
246 pub fn eq(&self, other: &FheUint8) -> FheBool {
247 BoolEvaluator::with_local_mut(|e| {
248 let key = ServerKeyEvaluationDomain::global();
249 arbitrary_bit_equality(e, self.data(), other.data(), key)
250 })
251 }
252
253 /// a != b
254 pub fn neq(&self, other: &FheUint8) -> FheBool {
255 BoolEvaluator::with_local_mut(|e| {
256 let key = ServerKeyEvaluationDomain::global();
257 let mut is_equal = arbitrary_bit_equality(e, self.data(), other.data(), key);
258 e.not_inplace(&mut is_equal);
259 is_equal
260 })
261 }
262
263 /// a < b
264 pub fn lt(&self, other: &FheUint8) -> FheBool {
265 BoolEvaluator::with_local_mut(|e| {
266 let key = ServerKeyEvaluationDomain::global();
267 arbitrary_bit_comparator(e, other.data(), self.data(), key)
268 })
269 }
270
271 /// a > b
272 pub fn gt(&self, other: &FheUint8) -> FheBool {
273 BoolEvaluator::with_local_mut(|e| {
274 let key = ServerKeyEvaluationDomain::global();
275 arbitrary_bit_comparator(e, self.data(), other.data(), key)
276 })
277 }
278
279 /// a <= b
280 pub fn le(&self, other: &FheUint8) -> FheBool {
281 BoolEvaluator::with_local_mut(|e| {
282 let key = ServerKeyEvaluationDomain::global();
283 let mut a_greater_b =
284 arbitrary_bit_comparator(e, self.data(), other.data(), key);
285 e.not_inplace(&mut a_greater_b);
286 a_greater_b
287 })
288 }
289
290 /// a >= b
291 pub fn ge(&self, other: &FheUint8) -> FheBool {
292 BoolEvaluator::with_local_mut(|e| {
293 let key = ServerKeyEvaluationDomain::global();
294 let mut a_less_b = arbitrary_bit_comparator(e, other.data(), self.data(), key);
295 e.not_inplace(&mut a_less_b);
296 a_less_b
297 })
298 }
299 }
300 }
301}
302
303#[cfg(test)]
304mod tests {
305 use itertools::Itertools;
306 use num_traits::Euclid;
307
308 use crate::{
309 bool::{
310 evaluator::{
311 aggregate_public_key_shares, aggregate_server_key_shares, gen_client_key, gen_keys,
312 gen_mp_keys_phase1, gen_mp_keys_phase2, set_mp_seed, set_parameter_set,
313 BoolEvaluator, ClientKey,
314 },
315 parameters::{MP_BOOL_PARAMS, SP_BOOL_PARAMS},
316 },
317 shortint::types::FheUint8,
318 Decryptor, Encryptor, MultiPartyDecryptor,
319 };
320
321 #[test]
322 fn all_uint8_apis() {
323 set_parameter_set(&SP_BOOL_PARAMS);
324
325 let (ck, sk) = gen_keys();
326 sk.set_server_key();
327
328 for i in 144..=255 {
329 for j in 100..=255 {
330 let m0 = i;
331 let m1 = j;
332 let c0 = ck.encrypt(&m0);
333 let c1 = ck.encrypt(&m1);
334
335 assert!(ck.decrypt(&c0) == m0);
336 assert!(ck.decrypt(&c1) == m1);
337
338 // Arithmetic
339 {
340 {
341 // Add
342 let mut c_m0_plus_m1 = FheUint8 {
343 data: c0.data().to_vec(),
344 };
345 c_m0_plus_m1 += &c1;
346 let m0_plus_m1 = ck.decrypt(&c_m0_plus_m1);
347 assert_eq!(
348 m0_plus_m1,
349 m0.wrapping_add(m1),
350 "Expected {} but got {m0_plus_m1} for {i}+{j}",
351 m0.wrapping_add(m1)
352 );
353 }
354 {
355 // Sub
356 let c_sub = &c0 - &c1;
357 let m0_sub_m1 = ck.decrypt(&c_sub);
358 assert_eq!(
359 m0_sub_m1,
360 m0.wrapping_sub(m1),
361 "Expected {} but got {m0_sub_m1} for {i}-{j}",
362 m0.wrapping_sub(m1)
363 );
364 }
365
366 {
367 // Mul
368 let c_m0m1 = &c0 * &c1;
369 let m0m1 = ck.decrypt(&c_m0m1);
370 assert_eq!(
371 m0m1,
372 m0.wrapping_mul(m1),
373 "Expected {} but got {m0m1} for {i}x{j}",
374 m0.wrapping_mul(m1)
375 );
376 }
377
378 // Div & Rem
379 {
380 let (c_quotient, c_rem) = c0.div_rem(&c1);
381 let m_quotient = ck.decrypt(&c_quotient);
382 let m_remainder = ck.decrypt(&c_rem);
383 if j != 0 {
384 let (q, r) = i.div_rem_euclid(&j);
385 assert_eq!(
386 m_quotient, q,
387 "Expected {} but got {m_quotient} for {i}/{j}",
388 q
389 );
390 assert_eq!(
391 m_remainder, r,
392 "Expected {} but got {m_quotient} for {i}%{j}",
393 r
394 );
395 } else {
396 assert_eq!(
397 m_quotient, 255,
398 "Expected 255 but got {m_quotient}. Case div by zero"
399 );
400 assert_eq!(
401 m_remainder, i,
402 "Expected {i} but got {m_quotient}. Case div by zero"
403 )
404 }
405 }
406 }
407
408 // Comparisons
409 {
410 {
411 let c_eq = c0.eq(&c1);
412 let is_eq = ck.decrypt(&c_eq);
413 assert_eq!(
414 is_eq,
415 i == j,
416 "Expected {} but got {is_eq} for {i}=={j}",
417 i == j
418 );
419 }
420
421 {
422 let c_gt = c0.gt(&c1);
423 let is_gt = ck.decrypt(&c_gt);
424 assert_eq!(
425 is_gt,
426 i > j,
427 "Expected {} but got {is_gt} for {i}>{j}",
428 i > j
429 );
430 }
431
432 {
433 let c_lt = c0.lt(&c1);
434 let is_lt = ck.decrypt(&c_lt);
435 assert_eq!(
436 is_lt,
437 i < j,
438 "Expected {} but got {is_lt} for {i}<{j}",
439 i < j
440 );
441 }
442
443 {
444 let c_ge = c0.ge(&c1);
445 let is_ge = ck.decrypt(&c_ge);
446 assert_eq!(
447 is_ge,
448 i >= j,
449 "Expected {} but got {is_ge} for {i}>={j}",
450 i >= j
451 );
452 }
453
454 {
455 let c_le = c0.le(&c1);
456 let is_le = ck.decrypt(&c_le);
457 assert_eq!(
458 is_le,
459 i <= j,
460 "Expected {} but got {is_le} for {i}<={j}",
461 i <= j
462 );
463 }
464 }
465 }
466 }
467 }
468
469 #[test]
470 fn fheuint8_test_multi_party() {
471 set_parameter_set(&MP_BOOL_PARAMS);
472 set_mp_seed([0; 32]);
473
474 let parties = 8;
475
476 // client keys and public key share
477 let cks = (0..parties)
478 .into_iter()
479 .map(|i| gen_client_key())
480 .collect_vec();
481
482 // round 1: generate pulic key shares
483 let pk_shares = cks.iter().map(|key| gen_mp_keys_phase1(key)).collect_vec();
484
485 let public_key = aggregate_public_key_shares(&pk_shares);
486
487 // round 2: generate server key shares
488 let server_key_shares = cks
489 .iter()
490 .map(|key| gen_mp_keys_phase2(key, &public_key))
491 .collect_vec();
492
493 // server aggregates the server key
494 let server_key = aggregate_server_key_shares(&server_key_shares);
495 server_key.set_server_key();
496
497 // Clients use Pk to encrypt private inputs
498 let a = 8u8;
499 let b = 10u8;
500 let c = 155u8;
501 let ct_a = public_key.encrypt(&a);
502 let ct_b = public_key.encrypt(&b);
503 let ct_c = public_key.encrypt(&c);
504
505 // server computes
506 // a*b + c
507 let mut ct_ab = &ct_a * &ct_b;
508 ct_ab += &ct_c;
509
510 // decrypt ab and check
511 // generate decryption shares
512 let decryption_shares = cks
513 .iter()
514 .map(|k| k.gen_decryption_share(&ct_ab))
515 .collect_vec();
516
517 // aggregate and decryption ab
518 let ab_add_c = cks[0].aggregate_decryption_shares(&ct_ab, &decryption_shares);
519 assert!(ab_add_c == (a.wrapping_mul(b)).wrapping_add(c));
520 }
521}