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