arnaucube
arnaucube/phantom-zonePublic Self-hosted repository
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. To do so call `reset_error_flags()`.
 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            /// Calculates `Self += rhs` and returns `overflow`
141            ///
142            /// `overflow` is set to `True` if `Self += rhs` overflowed,
143            /// otherwise it is set to `False`
144            pub fn overflowing_add_assign(&mut self, rhs: &FheUint8) -> FheBool {
145                BoolEvaluator::with_local_mut_mut(&mut |e| {
146                    let key = RuntimeServerKey::global();
147                    let (overflow, _) =
148                        arbitrary_bit_adder(e, self.data_mut(), rhs.data(), false, key);
149                    FheBool { data: overflow }
150                })
151            }
152
153            /// Returns (Self + rhs, overflow).
154            ///
155            /// `overflow` is set to `True` if `Self + rhs` overflowed,
156            /// otherwise it is set to `False`
157            pub fn overflowing_add(self, rhs: &FheUint8) -> (FheUint8, FheBool) {
158                BoolEvaluator::with_local_mut(|e| {
159                    let mut lhs = self.clone();
160                    let key = RuntimeServerKey::global();
161                    let (overflow, _) =
162                        arbitrary_bit_adder(e, lhs.data_mut(), rhs.data(), false, key);
163                    (lhs, FheBool { data: overflow })
164                })
165            }
166
167            /// Returns (Self - rhs, overflow).
168            ///
169            /// `overflow` is set to `True` if `Self - rhs` overflowed,
170            /// otherwise it is set to `False`
171            pub fn overflowing_sub(&self, rhs: &FheUint8) -> (FheUint8, FheBool) {
172                BoolEvaluator::with_local_mut(|e| {
173                    let key = RuntimeServerKey::global();
174                    let (out, mut overflow, _) =
175                        arbitrary_bit_subtractor(e, self.data(), rhs.data(), key);
176                    e.not_inplace(&mut overflow);
177                    (FheUint8 { data: out }, FheBool { data: overflow })
178                })
179            }
180
181            /// Returns (quotient, remainder) s.t. self = rhs x quotient +
182            /// remainder.
183            ///
184            /// If rhs is 0, then quotient = 255, remainder = self, and Div by
185            /// Zero error flag (accessible via `div_zero_error_flag`) is set to
186            /// `True`
187            pub fn div_rem(&self, rhs: &FheUint8) -> (FheUint8, FheUint8) {
188                // set div by 0 error flag
189                set_div_by_zero_flag(rhs);
190
191                BoolEvaluator::with_local_mut(|e| {
192                    let key = RuntimeServerKey::global();
193
194                    let (quotient, remainder) = arbitrary_bit_division_for_quotient_and_rem(
195                        e,
196                        self.data(),
197                        rhs.data(),
198                        key,
199                    );
200                    (FheUint8 { data: quotient }, FheUint8 { data: remainder })
201                })
202            }
203        }
204    }
205
206    mod booleans {
207        use crate::shortint::ops::{
208            arbitrary_bit_comparator, arbitrary_bit_equality, arbitrary_bit_mux,
209        };
210
211        use super::*;
212
213        impl FheUint8 {
214            /// Returns `FheBool` indicating `Self == other`
215            pub fn eq(&self, other: &FheUint8) -> FheBool {
216                BoolEvaluator::with_local_mut(|e| {
217                    let key = RuntimeServerKey::global();
218                    let out = arbitrary_bit_equality(e, self.data(), other.data(), key);
219                    FheBool { data: out }
220                })
221            }
222
223            /// Returns `FheBool` indicating `Self != other`
224            pub fn neq(&self, other: &FheUint8) -> FheBool {
225                BoolEvaluator::with_local_mut(|e| {
226                    let key = RuntimeServerKey::global();
227                    let mut is_equal = arbitrary_bit_equality(e, self.data(), other.data(), key);
228                    e.not_inplace(&mut is_equal);
229                    FheBool { data: is_equal }
230                })
231            }
232
233            /// Returns `FheBool` indicating `Self < other`
234            pub fn lt(&self, other: &FheUint8) -> FheBool {
235                BoolEvaluator::with_local_mut(|e| {
236                    let key = RuntimeServerKey::global();
237                    let out = arbitrary_bit_comparator(e, other.data(), self.data(), key);
238                    FheBool { data: out }
239                })
240            }
241
242            /// Returns `FheBool` indicating `Self > other`
243            pub fn gt(&self, other: &FheUint8) -> FheBool {
244                BoolEvaluator::with_local_mut(|e| {
245                    let key = RuntimeServerKey::global();
246                    let out = arbitrary_bit_comparator(e, self.data(), other.data(), key);
247                    FheBool { data: out }
248                })
249            }
250
251            /// Returns `FheBool` indicating `Self <= other`
252            pub fn le(&self, other: &FheUint8) -> FheBool {
253                BoolEvaluator::with_local_mut(|e| {
254                    let key = RuntimeServerKey::global();
255                    let mut a_greater_b =
256                        arbitrary_bit_comparator(e, self.data(), other.data(), key);
257                    e.not_inplace(&mut a_greater_b);
258                    FheBool { data: a_greater_b }
259                })
260            }
261
262            /// Returns `FheBool` indicating `Self >= other`
263            pub fn ge(&self, other: &FheUint8) -> FheBool {
264                BoolEvaluator::with_local_mut(|e| {
265                    let key = RuntimeServerKey::global();
266                    let mut a_less_b = arbitrary_bit_comparator(e, other.data(), self.data(), key);
267                    e.not_inplace(&mut a_less_b);
268                    FheBool { data: a_less_b }
269                })
270            }
271
272            /// Returns `Self` if `selector = True` else returns `other`
273            pub fn mux(&self, other: &FheUint8, selector: &FheBool) -> FheUint8 {
274                BoolEvaluator::with_local_mut(|e| {
275                    let key = RuntimeServerKey::global();
276                    let out = arbitrary_bit_mux(e, selector.data(), self.data(), other.data(), key);
277                    FheUint8 { data: out }
278                })
279            }
280
281            /// Returns max(`Self`, `other`)
282            pub fn max(&self, other: &FheUint8) -> FheUint8 {
283                let self_gt = self.gt(other);
284                self.mux(other, &self_gt)
285            }
286
287            /// Returns min(`Self`, `other`)
288            pub fn min(&self, other: &FheUint8) -> FheUint8 {
289                let self_lt = self.lt(other);
290                self.mux(other, &self_lt)
291            }
292        }
293    }
294}