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