phantom-zone/src/shortint/mod.rs
src/shortint/mod.rsRaw
1mod enc_dec;
2mod ops;
3
4pub type FheUint8 = enc_dec::FheUint8<Vec<u64>>;
5
6use std::cell::RefCell;
7
8use crate::bool::{BoolEvaluator, BooleanGates, FheBool, RuntimeServerKey};
9
10thread_local! {
11 static DIV_ZERO_ERROR: RefCell<Option<FheBool>> = RefCell::new(None);
12}
13
14/// Returns Boolean ciphertext indicating whether last division was attempeted
15/// with decnomiantor set to 0.
16pub fn div_zero_error_flag() -> Option<FheBool> {
17 DIV_ZERO_ERROR.with_borrow(|c| c.clone())
18}
19
20/// Reset all error flags
21///
22/// Error flags are thread local. When running multiple circuits in sequence
23/// within a single program you must prevent error flags set during the
24/// execution of previous circuit to affect error flags set during execution of
25/// the next circuit by resetting the flags before starting with next circuit.
26pub fn reset_error_flags() {
27 DIV_ZERO_ERROR.with_borrow_mut(|c| *c = None);
28}
29
30mod frontend {
31 use super::ops::{
32 arbitrary_bit_adder, arbitrary_bit_division_for_quotient_and_rem, arbitrary_bit_subtractor,
33 eight_bit_mul, is_zero,
34 };
35 use crate::utils::{Global, WithLocal};
36
37 use super::*;
38
39 /// Set Div by Zero flag after each divison. Div by zero flag is set to true
40 /// if either 1 of the division executed in circuit evaluation has
41 /// denominator set to 0.
42 fn set_div_by_zero_flag(denominator: &FheUint8) {
43 {
44 BoolEvaluator::with_local_mut(|e| {
45 let key = RuntimeServerKey::global();
46 let is_zero = is_zero(e, denominator.data(), key);
47 DIV_ZERO_ERROR.with_borrow_mut(|before_is_zero| {
48 if before_is_zero.is_none() {
49 *before_is_zero = Some(FheBool { data: is_zero });
50 } else {
51 e.or_inplace(before_is_zero.as_mut().unwrap().data_mut(), &is_zero, key);
52 }
53 });
54 })
55 }
56 }
57
58 mod arithetic {
59
60 use super::*;
61 use std::ops::{Add, AddAssign, Div, Mul, Rem, Sub};
62
63 impl AddAssign<&FheUint8> for FheUint8 {
64 fn add_assign(&mut self, rhs: &FheUint8) {
65 BoolEvaluator::with_local_mut_mut(&mut |e| {
66 let key = RuntimeServerKey::global();
67 arbitrary_bit_adder(e, self.data_mut(), rhs.data(), false, key);
68 });
69 }
70 }
71
72 impl Add<&FheUint8> for &FheUint8 {
73 type Output = FheUint8;
74 fn add(self, rhs: &FheUint8) -> Self::Output {
75 let mut a = self.clone();
76 a += rhs;
77 a
78 }
79 }
80
81 impl Sub<&FheUint8> for &FheUint8 {
82 type Output = FheUint8;
83 fn sub(self, rhs: &FheUint8) -> Self::Output {
84 BoolEvaluator::with_local_mut(|e| {
85 let key = RuntimeServerKey::global();
86 let (out, _, _) = arbitrary_bit_subtractor(e, self.data(), rhs.data(), key);
87 FheUint8 { data: out }
88 })
89 }
90 }
91
92 impl Mul<&FheUint8> for &FheUint8 {
93 type Output = FheUint8;
94 fn mul(self, rhs: &FheUint8) -> Self::Output {
95 BoolEvaluator::with_local_mut(|e| {
96 let key = RuntimeServerKey::global();
97 let out = eight_bit_mul(e, self.data(), rhs.data(), key);
98 FheUint8 { data: out }
99 })
100 }
101 }
102
103 impl Div<&FheUint8> for &FheUint8 {
104 type Output = FheUint8;
105 fn div(self, rhs: &FheUint8) -> Self::Output {
106 // set div by 0 error flag
107 set_div_by_zero_flag(rhs);
108
109 BoolEvaluator::with_local_mut(|e| {
110 let key = RuntimeServerKey::global();
111
112 let (quotient, _) = arbitrary_bit_division_for_quotient_and_rem(
113 e,
114 self.data(),
115 rhs.data(),
116 key,
117 );
118 FheUint8 { data: quotient }
119 })
120 }
121 }
122
123 impl Rem<&FheUint8> for &FheUint8 {
124 type Output = FheUint8;
125 fn rem(self, rhs: &FheUint8) -> Self::Output {
126 BoolEvaluator::with_local_mut(|e| {
127 let key = RuntimeServerKey::global();
128 let (_, remainder) = arbitrary_bit_division_for_quotient_and_rem(
129 e,
130 self.data(),
131 rhs.data(),
132 key,
133 );
134 FheUint8 { data: remainder }
135 })
136 }
137 }
138
139 impl FheUint8 {
140 pub fn overflowing_add_assign(&mut self, rhs: &FheUint8) -> FheBool {
141 BoolEvaluator::with_local_mut_mut(&mut |e| {
142 let key = RuntimeServerKey::global();
143 let (overflow, _) =
144 arbitrary_bit_adder(e, self.data_mut(), rhs.data(), false, key);
145 FheBool { data: overflow }
146 })
147 }
148
149 pub fn overflowing_add(self, rhs: &FheUint8) -> (FheUint8, FheBool) {
150 BoolEvaluator::with_local_mut(|e| {
151 let mut lhs = self.clone();
152 let key = RuntimeServerKey::global();
153 let (overflow, _) =
154 arbitrary_bit_adder(e, lhs.data_mut(), rhs.data(), false, key);
155 (lhs, FheBool { data: overflow })
156 })
157 }
158
159 pub fn overflowing_sub(&self, rhs: &FheUint8) -> (FheUint8, FheBool) {
160 BoolEvaluator::with_local_mut(|e| {
161 let key = RuntimeServerKey::global();
162 let (out, mut overflow, _) =
163 arbitrary_bit_subtractor(e, self.data(), rhs.data(), key);
164 e.not_inplace(&mut overflow);
165 (FheUint8 { data: out }, FheBool { data: overflow })
166 })
167 }
168
169 pub fn div_rem(&self, rhs: &FheUint8) -> (FheUint8, FheUint8) {
170 // set div by 0 error flag
171 set_div_by_zero_flag(rhs);
172
173 BoolEvaluator::with_local_mut(|e| {
174 let key = RuntimeServerKey::global();
175
176 let (quotient, remainder) = arbitrary_bit_division_for_quotient_and_rem(
177 e,
178 self.data(),
179 rhs.data(),
180 key,
181 );
182 (FheUint8 { data: quotient }, FheUint8 { data: remainder })
183 })
184 }
185 }
186 }
187
188 mod booleans {
189 use crate::shortint::ops::{
190 arbitrary_bit_comparator, arbitrary_bit_equality, arbitrary_bit_mux,
191 };
192
193 use super::*;
194
195 impl FheUint8 {
196 /// a == b
197 pub fn eq(&self, other: &FheUint8) -> FheBool {
198 BoolEvaluator::with_local_mut(|e| {
199 let key = RuntimeServerKey::global();
200 let out = arbitrary_bit_equality(e, self.data(), other.data(), key);
201 FheBool { data: out }
202 })
203 }
204
205 /// a != b
206 pub fn neq(&self, other: &FheUint8) -> FheBool {
207 BoolEvaluator::with_local_mut(|e| {
208 let key = RuntimeServerKey::global();
209 let mut is_equal = arbitrary_bit_equality(e, self.data(), other.data(), key);
210 e.not_inplace(&mut is_equal);
211 FheBool { data: is_equal }
212 })
213 }
214
215 /// a < b
216 pub fn lt(&self, other: &FheUint8) -> FheBool {
217 BoolEvaluator::with_local_mut(|e| {
218 let key = RuntimeServerKey::global();
219 let out = arbitrary_bit_comparator(e, other.data(), self.data(), key);
220 FheBool { data: out }
221 })
222 }
223
224 /// a > b
225 pub fn gt(&self, other: &FheUint8) -> FheBool {
226 BoolEvaluator::with_local_mut(|e| {
227 let key = RuntimeServerKey::global();
228 let out = arbitrary_bit_comparator(e, self.data(), other.data(), key);
229 FheBool { data: out }
230 })
231 }
232
233 /// a <= b
234 pub fn le(&self, other: &FheUint8) -> FheBool {
235 BoolEvaluator::with_local_mut(|e| {
236 let key = RuntimeServerKey::global();
237 let mut a_greater_b =
238 arbitrary_bit_comparator(e, self.data(), other.data(), key);
239 e.not_inplace(&mut a_greater_b);
240 FheBool { data: a_greater_b }
241 })
242 }
243
244 /// a >= b
245 pub fn ge(&self, other: &FheUint8) -> FheBool {
246 BoolEvaluator::with_local_mut(|e| {
247 let key = RuntimeServerKey::global();
248 let mut a_less_b = arbitrary_bit_comparator(e, other.data(), self.data(), key);
249 e.not_inplace(&mut a_less_b);
250 FheBool { data: a_less_b }
251 })
252 }
253
254 /// Returns `Self` if `selector = True` else returns `other`
255 pub fn mux(&self, other: &FheUint8, selector: &FheBool) -> FheUint8 {
256 BoolEvaluator::with_local_mut(|e| {
257 let key = RuntimeServerKey::global();
258 let out = arbitrary_bit_mux(e, selector.data(), self.data(), other.data(), key);
259 FheUint8 { data: out }
260 })
261 }
262
263 /// max(`Self`, `other`)
264 pub fn max(&self, other: &FheUint8) -> FheUint8 {
265 let self_gt = self.gt(other);
266 self.mux(other, &self_gt)
267 }
268
269 /// min(`Self`, `other`)
270 pub fn min(&self, other: &FheUint8) -> FheUint8 {
271 let self_lt = self.lt(other);
272 self.mux(other, &self_lt)
273 }
274 }
275 }
276}