phantom-zone/src/shortint/mod.rs
src/shortint/mod.rsRaw
1use itertools::Itertools;
2
3use crate::{
4 bool::evaluator::{BoolEvaluator, ClientKey, ServerKeyEvaluationDomain, BOOL_SERVER_KEY},
5 utils::{Global, WithLocal},
6 Decryptor, Encryptor,
7};
8
9mod ops;
10mod types;
11
12type FheUint8 = types::FheUint8<Vec<u64>>;
13
14impl Encryptor<u8, FheUint8> for ClientKey {
15 fn encrypt(&self, m: &u8) -> FheUint8 {
16 let cts = (0..8)
17 .into_iter()
18 .map(|i| {
19 let bit = ((m >> i) & 1) == 1;
20 Encryptor::<bool, Vec<u64>>::encrypt(self, &bit)
21 })
22 .collect_vec();
23 FheUint8 { data: cts }
24 }
25}
26
27impl Decryptor<u8, FheUint8> for ClientKey {
28 fn decrypt(&self, c: &FheUint8) -> u8 {
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
40mod frontend {
41 use super::ops::{
42 arbitrary_bit_adder, arbitrary_bit_division_for_quotient_and_rem, arbitrary_bit_subtractor,
43 eight_bit_mul,
44 };
45 use crate::{
46 bool::evaluator::{BoolEvaluator, ServerKeyEvaluationDomain},
47 utils::{Global, WithLocal},
48 };
49
50 use super::FheUint8;
51
52 mod arithetic {
53 use crate::bool::{evaluator::BooleanGates, FheBool};
54
55 use super::*;
56 use std::ops::{Add, AddAssign, Div, Mul, Rem, Sub};
57
58 impl AddAssign<&FheUint8> for FheUint8 {
59 fn add_assign(&mut self, rhs: &FheUint8) {
60 BoolEvaluator::with_local_mut_mut(&mut |e| {
61 let key = ServerKeyEvaluationDomain::global();
62 arbitrary_bit_adder(e, self.data_mut(), rhs.data(), false, key);
63 });
64 }
65 }
66
67 impl Add<&FheUint8> for &FheUint8 {
68 type Output = FheUint8;
69 fn add(self, rhs: &FheUint8) -> Self::Output {
70 let mut a = self.clone();
71 a += rhs;
72 a
73 }
74 }
75
76 impl Sub<&FheUint8> for &FheUint8 {
77 type Output = FheUint8;
78 fn sub(self, rhs: &FheUint8) -> Self::Output {
79 BoolEvaluator::with_local_mut(|e| {
80 let key = ServerKeyEvaluationDomain::global();
81 let (out, _, _) = arbitrary_bit_subtractor(e, self.data(), rhs.data(), key);
82 FheUint8 { data: out }
83 })
84 }
85 }
86
87 impl Mul<&FheUint8> for &FheUint8 {
88 type Output = FheUint8;
89 fn mul(self, rhs: &FheUint8) -> Self::Output {
90 BoolEvaluator::with_local_mut(|e| {
91 let key = ServerKeyEvaluationDomain::global();
92 let out = eight_bit_mul(e, self.data(), rhs.data(), key);
93 FheUint8 { data: out }
94 })
95 }
96 }
97
98 impl Div<&FheUint8> for &FheUint8 {
99 type Output = FheUint8;
100 fn div(self, rhs: &FheUint8) -> Self::Output {
101 // TODO(Jay:) Figure out how to set zero error flag
102 BoolEvaluator::with_local_mut(|e| {
103 let key = ServerKeyEvaluationDomain::global();
104 let (quotient, _) = arbitrary_bit_division_for_quotient_and_rem(
105 e,
106 self.data(),
107 rhs.data(),
108 key,
109 );
110 FheUint8 { data: quotient }
111 })
112 }
113 }
114
115 impl Rem<&FheUint8> for &FheUint8 {
116 type Output = FheUint8;
117 fn rem(self, rhs: &FheUint8) -> Self::Output {
118 BoolEvaluator::with_local_mut(|e| {
119 let key = ServerKeyEvaluationDomain::global();
120 let (_, remainder) = arbitrary_bit_division_for_quotient_and_rem(
121 e,
122 self.data(),
123 rhs.data(),
124 key,
125 );
126 FheUint8 { data: remainder }
127 })
128 }
129 }
130
131 impl FheUint8 {
132 pub fn overflowing_add_assign(&mut self, rhs: &FheUint8) -> FheBool {
133 BoolEvaluator::with_local_mut_mut(&mut |e| {
134 let key = ServerKeyEvaluationDomain::global();
135 let (overflow, _) =
136 arbitrary_bit_adder(e, self.data_mut(), rhs.data(), false, key);
137 overflow
138 })
139 }
140
141 pub fn overflowing_add(self, rhs: &FheUint8) -> (FheUint8, FheBool) {
142 BoolEvaluator::with_local_mut(|e| {
143 let mut lhs = self.clone();
144 let key = ServerKeyEvaluationDomain::global();
145 let (overflow, _) =
146 arbitrary_bit_adder(e, lhs.data_mut(), rhs.data(), false, key);
147 (lhs, overflow)
148 })
149 }
150
151 pub fn overflowing_sub(&self, rhs: &FheUint8) -> (FheUint8, FheBool) {
152 BoolEvaluator::with_local_mut(|e| {
153 let key = ServerKeyEvaluationDomain::global();
154 let (out, mut overflow, _) =
155 arbitrary_bit_subtractor(e, self.data(), rhs.data(), key);
156 e.not_inplace(&mut overflow);
157 (FheUint8 { data: out }, overflow)
158 })
159 }
160
161 pub fn div_rem(&self, rhs: &FheUint8) -> (FheUint8, FheUint8) {
162 // TODO(Jay:) Figure out how to set zero error flag
163 BoolEvaluator::with_local_mut(|e| {
164 let key = ServerKeyEvaluationDomain::global();
165 let (quotient, remainder) = arbitrary_bit_division_for_quotient_and_rem(
166 e,
167 self.data(),
168 rhs.data(),
169 key,
170 );
171 (FheUint8 { data: quotient }, FheUint8 { data: remainder })
172 })
173 }
174 }
175 }
176
177 mod booleans {
178 use crate::{
179 bool::{evaluator::BooleanGates, FheBool},
180 shortint::ops::{
181 arbitrary_bit_comparator, arbitrary_bit_equality, arbitrary_signed_bit_comparator,
182 },
183 };
184
185 use super::*;
186
187 impl FheUint8 {
188 /// a == b
189 pub fn eq(&self, other: &FheUint8) -> FheBool {
190 BoolEvaluator::with_local_mut(|e| {
191 let key = ServerKeyEvaluationDomain::global();
192 arbitrary_bit_equality(e, self.data(), other.data(), key)
193 })
194 }
195
196 /// a != b
197 pub fn neq(&self, other: &FheUint8) -> FheBool {
198 BoolEvaluator::with_local_mut(|e| {
199 let key = ServerKeyEvaluationDomain::global();
200 let mut is_equal = arbitrary_bit_equality(e, self.data(), other.data(), key);
201 e.not_inplace(&mut is_equal);
202 is_equal
203 })
204 }
205
206 /// a < b
207 pub fn lt(&self, other: &FheUint8) -> FheBool {
208 BoolEvaluator::with_local_mut(|e| {
209 let key = ServerKeyEvaluationDomain::global();
210 arbitrary_bit_comparator(e, other.data(), self.data(), key)
211 })
212 }
213
214 /// a > b
215 pub fn gt(&self, other: &FheUint8) -> FheBool {
216 BoolEvaluator::with_local_mut(|e| {
217 let key = ServerKeyEvaluationDomain::global();
218 arbitrary_bit_comparator(e, self.data(), other.data(), key)
219 })
220 }
221
222 /// a <= b
223 pub fn le(&self, other: &FheUint8) -> FheBool {
224 BoolEvaluator::with_local_mut(|e| {
225 let key = ServerKeyEvaluationDomain::global();
226 let mut a_greater_b =
227 arbitrary_bit_comparator(e, self.data(), other.data(), key);
228 e.not_inplace(&mut a_greater_b);
229 a_greater_b
230 })
231 }
232
233 /// a >= b
234 pub fn ge(&self, other: &FheUint8) -> FheBool {
235 BoolEvaluator::with_local_mut(|e| {
236 let key = ServerKeyEvaluationDomain::global();
237 let mut a_less_b = arbitrary_bit_comparator(e, other.data(), self.data(), key);
238 e.not_inplace(&mut a_less_b);
239 a_less_b
240 })
241 }
242 }
243 }
244}
245
246#[cfg(test)]
247mod tests {
248 use num_traits::Euclid;
249
250 use crate::{
251 bool::{
252 evaluator::{gen_keys, set_parameter_set, BoolEvaluator},
253 parameters::SP_BOOL_PARAMS,
254 },
255 shortint::types::FheUint8,
256 Decryptor, Encryptor,
257 };
258
259 #[test]
260 fn all_uint8_apis() {
261 set_parameter_set(&SP_BOOL_PARAMS);
262
263 let (ck, sk) = gen_keys();
264 sk.set_server_key();
265
266 for i in 144..=255 {
267 for j in 100..=255 {
268 let m0 = i;
269 let m1 = j;
270 let c0 = ck.encrypt(&m0);
271 let c1 = ck.encrypt(&m1);
272
273 assert!(ck.decrypt(&c0) == m0);
274 assert!(ck.decrypt(&c1) == m1);
275
276 // Arithmetic
277 {
278 {
279 // Add
280 let mut c_m0_plus_m1 = FheUint8 {
281 data: c0.data().to_vec(),
282 };
283 c_m0_plus_m1 += &c1;
284 let m0_plus_m1 = ck.decrypt(&c_m0_plus_m1);
285 assert_eq!(
286 m0_plus_m1,
287 m0.wrapping_add(m1),
288 "Expected {} but got {m0_plus_m1} for {i}+{j}",
289 m0.wrapping_add(m1)
290 );
291 }
292 {
293 // Sub
294 let c_sub = &c0 - &c1;
295 let m0_sub_m1 = ck.decrypt(&c_sub);
296 assert_eq!(
297 m0_sub_m1,
298 m0.wrapping_sub(m1),
299 "Expected {} but got {m0_sub_m1} for {i}-{j}",
300 m0.wrapping_sub(m1)
301 );
302 }
303
304 {
305 // Mul
306 let c_m0m1 = &c0 * &c1;
307 let m0m1 = ck.decrypt(&c_m0m1);
308 assert_eq!(
309 m0m1,
310 m0.wrapping_mul(m1),
311 "Expected {} but got {m0m1} for {i}x{j}",
312 m0.wrapping_mul(m1)
313 );
314 }
315
316 // Div & Rem
317 {
318 let (c_quotient, c_rem) = c0.div_rem(&c1);
319 let m_quotient = ck.decrypt(&c_quotient);
320 let m_remainder = ck.decrypt(&c_rem);
321 if j != 0 {
322 let (q, r) = i.div_rem_euclid(&j);
323 assert_eq!(
324 m_quotient, q,
325 "Expected {} but got {m_quotient} for {i}/{j}",
326 q
327 );
328 assert_eq!(
329 m_remainder, r,
330 "Expected {} but got {m_quotient} for {i}%{j}",
331 r
332 );
333 } else {
334 assert_eq!(
335 m_quotient, 255,
336 "Expected 255 but got {m_quotient}. Case div by zero"
337 );
338 assert_eq!(
339 m_remainder, i,
340 "Expected {i} but got {m_quotient}. Case div by zero"
341 )
342 }
343 }
344 }
345
346 // Comparisons
347 {
348 {
349 let c_eq = c0.eq(&c1);
350 let is_eq = ck.decrypt(&c_eq);
351 assert_eq!(
352 is_eq,
353 i == j,
354 "Expected {} but got {is_eq} for {i}=={j}",
355 i == j
356 );
357 }
358
359 {
360 let c_gt = c0.gt(&c1);
361 let is_gt = ck.decrypt(&c_gt);
362 assert_eq!(
363 is_gt,
364 i > j,
365 "Expected {} but got {is_gt} for {i}>{j}",
366 i > j
367 );
368 }
369
370 {
371 let c_lt = c0.lt(&c1);
372 let is_lt = ck.decrypt(&c_lt);
373 assert_eq!(
374 is_lt,
375 i < j,
376 "Expected {} but got {is_lt} for {i}<{j}",
377 i < j
378 );
379 }
380
381 {
382 let c_ge = c0.ge(&c1);
383 let is_ge = ck.decrypt(&c_ge);
384 assert_eq!(
385 is_ge,
386 i >= j,
387 "Expected {} but got {is_ge} for {i}>={j}",
388 i >= j
389 );
390 }
391
392 {
393 let c_le = c0.le(&c1);
394 let is_le = ck.decrypt(&c_le);
395 assert_eq!(
396 is_le,
397 i <= j,
398 "Expected {} but got {is_le} for {i}<={j}",
399 i <= j
400 );
401 }
402 }
403 }
404 }
405 }
406}