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