arnaucube
arnaucube/phantom-zonePublic Self-hosted repository
phantom-zone/src/shortint/mod.rs
src/shortint/mod.rsRaw
  1mod enc_dec;
  2mod ops;
  3mod types;
  4
  5pub type FheUint8 = enc_dec::FheUint8<Vec<u64>>;
  6pub type FheBool = Vec<u64>;
  7
  8use crate::bool::{evaluator::BooleanGates, BoolEvaluator, RuntimeServerKey};
  9
 10mod frontend {
 11    use super::ops::{
 12        arbitrary_bit_adder, arbitrary_bit_division_for_quotient_and_rem, arbitrary_bit_subtractor,
 13        eight_bit_mul,
 14    };
 15    use crate::utils::{Global, WithLocal};
 16
 17    use super::*;
 18
 19    mod arithetic {
 20
 21        use super::*;
 22        use std::ops::{Add, AddAssign, Div, Mul, Rem, Sub};
 23
 24        impl AddAssign<&FheUint8> for FheUint8 {
 25            fn add_assign(&mut self, rhs: &FheUint8) {
 26                BoolEvaluator::with_local_mut_mut(&mut |e| {
 27                    let key = RuntimeServerKey::global();
 28                    arbitrary_bit_adder(e, self.data_mut(), rhs.data(), false, key);
 29                });
 30            }
 31        }
 32
 33        impl Add<&FheUint8> for &FheUint8 {
 34            type Output = FheUint8;
 35            fn add(self, rhs: &FheUint8) -> Self::Output {
 36                let mut a = self.clone();
 37                a += rhs;
 38                a
 39            }
 40        }
 41
 42        impl Sub<&FheUint8> for &FheUint8 {
 43            type Output = FheUint8;
 44            fn sub(self, rhs: &FheUint8) -> Self::Output {
 45                BoolEvaluator::with_local_mut(|e| {
 46                    let key = RuntimeServerKey::global();
 47                    let (out, _, _) = arbitrary_bit_subtractor(e, self.data(), rhs.data(), key);
 48                    FheUint8 { data: out }
 49                })
 50            }
 51        }
 52
 53        impl Mul<&FheUint8> for &FheUint8 {
 54            type Output = FheUint8;
 55            fn mul(self, rhs: &FheUint8) -> Self::Output {
 56                BoolEvaluator::with_local_mut(|e| {
 57                    let key = RuntimeServerKey::global();
 58                    let out = eight_bit_mul(e, self.data(), rhs.data(), key);
 59                    FheUint8 { data: out }
 60                })
 61            }
 62        }
 63
 64        impl Div<&FheUint8> for &FheUint8 {
 65            type Output = FheUint8;
 66            fn div(self, rhs: &FheUint8) -> Self::Output {
 67                // TODO(Jay:) Figure out how to set zero error flag
 68                BoolEvaluator::with_local_mut(|e| {
 69                    let key = RuntimeServerKey::global();
 70                    let (quotient, _) = arbitrary_bit_division_for_quotient_and_rem(
 71                        e,
 72                        self.data(),
 73                        rhs.data(),
 74                        key,
 75                    );
 76                    FheUint8 { data: quotient }
 77                })
 78            }
 79        }
 80
 81        impl Rem<&FheUint8> for &FheUint8 {
 82            type Output = FheUint8;
 83            fn rem(self, rhs: &FheUint8) -> Self::Output {
 84                BoolEvaluator::with_local_mut(|e| {
 85                    let key = RuntimeServerKey::global();
 86                    let (_, remainder) = arbitrary_bit_division_for_quotient_and_rem(
 87                        e,
 88                        self.data(),
 89                        rhs.data(),
 90                        key,
 91                    );
 92                    FheUint8 { data: remainder }
 93                })
 94            }
 95        }
 96
 97        impl FheUint8 {
 98            pub fn overflowing_add_assign(&mut self, rhs: &FheUint8) -> FheBool {
 99                BoolEvaluator::with_local_mut_mut(&mut |e| {
100                    let key = RuntimeServerKey::global();
101                    let (overflow, _) =
102                        arbitrary_bit_adder(e, self.data_mut(), rhs.data(), false, key);
103                    overflow
104                })
105            }
106
107            pub fn overflowing_add(self, rhs: &FheUint8) -> (FheUint8, FheBool) {
108                BoolEvaluator::with_local_mut(|e| {
109                    let mut lhs = self.clone();
110                    let key = RuntimeServerKey::global();
111                    let (overflow, _) =
112                        arbitrary_bit_adder(e, lhs.data_mut(), rhs.data(), false, key);
113                    (lhs, overflow)
114                })
115            }
116
117            pub fn overflowing_sub(&self, rhs: &FheUint8) -> (FheUint8, FheBool) {
118                BoolEvaluator::with_local_mut(|e| {
119                    let key = RuntimeServerKey::global();
120                    let (out, mut overflow, _) =
121                        arbitrary_bit_subtractor(e, self.data(), rhs.data(), key);
122                    e.not_inplace(&mut overflow);
123                    (FheUint8 { data: out }, overflow)
124                })
125            }
126
127            pub fn div_rem(&self, rhs: &FheUint8) -> (FheUint8, FheUint8) {
128                // TODO(Jay:) Figure out how to set zero error flag
129                BoolEvaluator::with_local_mut(|e| {
130                    let key = RuntimeServerKey::global();
131                    let (quotient, remainder) = arbitrary_bit_division_for_quotient_and_rem(
132                        e,
133                        self.data(),
134                        rhs.data(),
135                        key,
136                    );
137                    (FheUint8 { data: quotient }, FheUint8 { data: remainder })
138                })
139            }
140        }
141    }
142
143    mod booleans {
144        use crate::shortint::ops::{
145            arbitrary_bit_comparator, arbitrary_bit_equality, arbitrary_signed_bit_comparator,
146        };
147
148        use super::*;
149
150        impl FheUint8 {
151            /// a == b
152            pub fn eq(&self, other: &FheUint8) -> FheBool {
153                BoolEvaluator::with_local_mut(|e| {
154                    let key = RuntimeServerKey::global();
155                    arbitrary_bit_equality(e, self.data(), other.data(), key)
156                })
157            }
158
159            /// a != b
160            pub fn neq(&self, other: &FheUint8) -> FheBool {
161                BoolEvaluator::with_local_mut(|e| {
162                    let key = RuntimeServerKey::global();
163                    let mut is_equal = arbitrary_bit_equality(e, self.data(), other.data(), key);
164                    e.not_inplace(&mut is_equal);
165                    is_equal
166                })
167            }
168
169            /// a < b
170            pub fn lt(&self, other: &FheUint8) -> FheBool {
171                BoolEvaluator::with_local_mut(|e| {
172                    let key = RuntimeServerKey::global();
173                    arbitrary_bit_comparator(e, other.data(), self.data(), key)
174                })
175            }
176
177            /// a > b
178            pub fn gt(&self, other: &FheUint8) -> FheBool {
179                BoolEvaluator::with_local_mut(|e| {
180                    let key = RuntimeServerKey::global();
181                    arbitrary_bit_comparator(e, self.data(), other.data(), key)
182                })
183            }
184
185            /// a <= b
186            pub fn le(&self, other: &FheUint8) -> FheBool {
187                BoolEvaluator::with_local_mut(|e| {
188                    let key = RuntimeServerKey::global();
189                    let mut a_greater_b =
190                        arbitrary_bit_comparator(e, self.data(), other.data(), key);
191                    e.not_inplace(&mut a_greater_b);
192                    a_greater_b
193                })
194            }
195
196            /// a >= b
197            pub fn ge(&self, other: &FheUint8) -> FheBool {
198                BoolEvaluator::with_local_mut(|e| {
199                    let key = RuntimeServerKey::global();
200                    let mut a_less_b = arbitrary_bit_comparator(e, other.data(), self.data(), key);
201                    e.not_inplace(&mut a_less_b);
202                    a_less_b
203                })
204            }
205        }
206    }
207}
208
209#[cfg(test)]
210mod tests {
211    use itertools::Itertools;
212    use num_traits::Euclid;
213
214    use crate::{
215        bool::set_parameter_set, shortint::enc_dec::FheUint8, utils::WithLocal, Decryptor,
216        Encryptor, MultiPartyDecryptor,
217    };
218
219    use super::*;
220}