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};
8use ops::{
9 arbitrary_bit_adder, arbitrary_bit_division_for_quotient_and_rem, arbitrary_bit_subtractor,
10 eight_bit_mul,
11};
12
13mod ops;
14mod types;
15
16type FheUint8 = types::FheUint8<Vec<u64>>;
17
18fn add_mut(a: &mut FheUint8, b: &FheUint8) {
19 BoolEvaluator::with_local_mut_mut(&mut |e| {
20 let key = ServerKeyEvaluationDomain::global();
21 arbitrary_bit_adder(e, a.data_mut(), b.data(), false, key);
22 });
23}
24
25fn sub(a: &FheUint8, b: &FheUint8) -> FheUint8 {
26 BoolEvaluator::with_local_mut(|e| {
27 let key = ServerKeyEvaluationDomain::global();
28 let (out, _, _) = arbitrary_bit_subtractor(e, a.data(), b.data(), key);
29 FheUint8 { data: out }
30 })
31}
32
33fn mul(a: &FheUint8, b: &FheUint8) -> FheUint8 {
34 BoolEvaluator::with_local_mut(|e| {
35 let key = ServerKeyEvaluationDomain::global();
36 let out = eight_bit_mul(e, a.data(), b.data(), key);
37 FheUint8 { data: out }
38 })
39}
40
41fn div(a: &FheUint8, b: &FheUint8) -> (FheUint8, FheUint8) {
42 BoolEvaluator::with_local_mut(|e| {
43 let key = ServerKeyEvaluationDomain::global();
44 let (quotient, remainder) =
45 arbitrary_bit_division_for_quotient_and_rem(e, a.data(), b.data(), key);
46
47 (FheUint8 { data: quotient }, FheUint8 { data: remainder })
48 })
49}
50
51impl Encryptor<u8, FheUint8> for ClientKey {
52 fn encrypt(&self, m: &u8) -> FheUint8 {
53 let cts = (0..8)
54 .into_iter()
55 .map(|i| {
56 let bit = ((m >> i) & 1) == 1;
57 Encryptor::<bool, Vec<u64>>::encrypt(self, &bit)
58 })
59 .collect_vec();
60 FheUint8 { data: cts }
61 }
62}
63
64impl Decryptor<u8, FheUint8> for ClientKey {
65 fn decrypt(&self, c: &FheUint8) -> u8 {
66 let mut out = 0u8;
67 c.data().iter().enumerate().for_each(|(index, bit_c)| {
68 let bool = Decryptor::<bool, Vec<u64>>::decrypt(self, bit_c);
69 if bool {
70 out += 1 << index;
71 }
72 });
73 out
74 }
75}
76
77mod frontend {
78 use super::ops::{
79 arbitrary_bit_adder, arbitrary_bit_division_for_quotient_and_rem, arbitrary_bit_subtractor,
80 eight_bit_mul,
81 };
82 use crate::{
83 bool::evaluator::{BoolEvaluator, ServerKeyEvaluationDomain},
84 utils::{Global, WithLocal},
85 };
86
87 use super::{add_mut, div, mul, FheUint8};
88
89 mod arithetic {
90 use super::*;
91 use std::ops::{Add, AddAssign, Div, Mul, Rem, Sub};
92
93 impl AddAssign<&FheUint8> for FheUint8 {
94 fn add_assign(&mut self, rhs: &FheUint8) {
95 BoolEvaluator::with_local_mut_mut(&mut |e| {
96 let key = ServerKeyEvaluationDomain::global();
97 arbitrary_bit_adder(e, self.data_mut(), rhs.data(), false, key);
98 });
99 }
100 }
101
102 impl Add<&FheUint8> for &FheUint8 {
103 type Output = FheUint8;
104 fn add(self, rhs: &FheUint8) -> Self::Output {
105 let mut a = self.clone();
106 a += rhs;
107 a
108 }
109 }
110
111 impl Sub<&FheUint8> for &FheUint8 {
112 type Output = FheUint8;
113 fn sub(self, rhs: &FheUint8) -> Self::Output {
114 BoolEvaluator::with_local_mut(|e| {
115 let key = ServerKeyEvaluationDomain::global();
116 let (out, _, _) = arbitrary_bit_subtractor(e, self.data(), self.data(), key);
117 FheUint8 { data: out }
118 })
119 }
120 }
121
122 impl Mul<&FheUint8> for &FheUint8 {
123 type Output = FheUint8;
124 fn mul(self, rhs: &FheUint8) -> Self::Output {
125 BoolEvaluator::with_local_mut(|e| {
126 let key = ServerKeyEvaluationDomain::global();
127 let out = eight_bit_mul(e, self.data(), rhs.data(), key);
128 FheUint8 { data: out }
129 })
130 }
131 }
132
133 impl Div<&FheUint8> for &FheUint8 {
134 type Output = FheUint8;
135 fn div(self, rhs: &FheUint8) -> Self::Output {
136 BoolEvaluator::with_local_mut(|e| {
137 let key = ServerKeyEvaluationDomain::global();
138 let (quotient, _) = arbitrary_bit_division_for_quotient_and_rem(
139 e,
140 self.data(),
141 rhs.data(),
142 key,
143 );
144 FheUint8 { data: quotient }
145 })
146 }
147 }
148
149 impl Rem<&FheUint8> for &FheUint8 {
150 type Output = FheUint8;
151 fn rem(self, rhs: &FheUint8) -> Self::Output {
152 BoolEvaluator::with_local_mut(|e| {
153 let key = ServerKeyEvaluationDomain::global();
154 let (_, remainder) = arbitrary_bit_division_for_quotient_and_rem(
155 e,
156 self.data(),
157 rhs.data(),
158 key,
159 );
160 FheUint8 { data: remainder }
161 })
162 }
163 }
164 }
165
166 mod booleans {}
167}
168
169#[cfg(test)]
170mod tests {
171 use num_traits::Euclid;
172
173 use crate::{
174 bool::{
175 evaluator::{gen_keys, set_parameter_set, BoolEvaluator},
176 parameters::SP_BOOL_PARAMS,
177 },
178 shortint::{add_mut, div, mul, sub, types::FheUint8},
179 Decryptor, Encryptor,
180 };
181
182 #[test]
183 fn qwerty() {
184 set_parameter_set(&SP_BOOL_PARAMS);
185
186 let (ck, sk) = gen_keys();
187 sk.set_server_key();
188
189 for i in 1..=255 {
190 for j in 0..=255 {
191 let m0 = i;
192 let m1 = j;
193 let c0 = ck.encrypt(&m0);
194 let c1 = ck.encrypt(&m1);
195
196 assert!(ck.decrypt(&c0) == m0);
197 assert!(ck.decrypt(&c1) == m1);
198
199 // Add
200 // let mut c_m0_plus_m1 = FheUint8 {
201 // data: c0.data().to_vec(),
202 // };
203 // add_mut(&mut c_m0_plus_m1, &c1);
204 // let m0_plus_m1 = ck.decrypt(&c_m0_plus_m1);
205 // assert_eq!(
206 // m0_plus_m1,
207 // m0.wrapping_add(m1),
208 // "Expected {} but got {m0_plus_m1} for {i}+{j}",
209 // m0.wrapping_add(m1)
210 // );
211
212 // Sub
213 // let c_sub = sub(&c0, &c1);
214 // let m0_sub_m1 = ck.decrypt(&c_sub);
215 // dbg!(m0, m1, m0_sub_m1);
216 // assert_eq!(
217 // m0_sub_m1,
218 // m0.wrapping_sub(m1),
219 // "Expected {} but got {m0_sub_m1} for {i}-{j}",
220 // m0.wrapping_sub(m1)
221 // );
222
223 // Mul
224 // let c_m0m1 = mul(&c0, &c1);
225 // let m0m1 = ck.decrypt(&c_m0m1);
226 // assert_eq!(
227 // m0m1,
228 // m0.wrapping_mul(m1),
229 // "Expected {} but got {m0m1} for {i}x{j}",
230 // m0.wrapping_mul(m1)
231 // );
232
233 // Div
234 // let (c_quotient, c_rem) = div(&c0, &c1);
235 // let m_quotient = ck.decrypt(&c_quotient);
236 // let m_remainder = ck.decrypt(&c_rem);
237 // if j != 0 {
238 // let (q, r) = i.div_rem_euclid(&j);
239 // assert_eq!(
240 // m_quotient, q,
241 // "Expected {} but got {m_quotient} for {i}/{j}",
242 // q
243 // );
244 // assert_eq!(
245 // m_remainder, r,
246 // "Expected {} but got {m_quotient} for {i}%{j}",
247 // r
248 // );
249 // } else {
250 // assert_eq!(
251 // m_quotient, 255,
252 // "Expected 255 but got {m_quotient}. Case div by zero"
253 // );
254 // assert_eq!(
255 // m_remainder, i,
256 // "Expected {i} but got {m_quotient}. Case div by zero"
257 // )
258 // }
259 }
260 }
261 }
262}