1use itertools::Itertools;
  2
  3use crate::{
  4    bool::BoolEvaluator,
  5    random::{DefaultSecureRng, RandomFillUniformInModulus},
  6    utils::WithLocal,
  7    Decryptor, Encryptor, KeySwitchWithId, Matrix, MatrixEntity, MatrixMut, MultiPartyDecryptor,
  8    RowMut, SampleExtractor,
  9};
 10
 11/// Fhe UInt8 type
 12///
 13/// - Stores encryptions of bits in little endian (i.e least signficant bit
 14///   stored at 0th index and most signficant bit stores at 7th index)
 15#[derive(Clone)]
 16pub struct FheUint8<C> {
 17    pub(super) data: Vec<C>,
 18}
 19
 20impl<C> FheUint8<C> {
 21    pub(super) fn data(&self) -> &[C] {
 22        &self.data
 23    }
 24
 25    pub(super) fn data_mut(&mut self) -> &mut [C] {
 26        &mut self.data
 27    }
 28}
 29
 30/// Stored a batch of Fhe Uint8 ciphertext as collection of RLWE ciphertexts
 31///
 32/// To extract Fhe Uint8 ciphertext at `index` call `self.extract(index)`
 33pub struct BatchedFheUint8<C> {
 34    /// Vector of RLWE ciphertexts `C`
 35    data: Vec<C>,
 36    /// Count of FheUint8s packed in vector of RLWE ciphertexts
 37    count: usize,
 38}
 39
 40impl<C, R> SampleExtractor<FheUint8<R>> for BatchedFheUint8<C>
 41where
 42    C: SampleExtractor<R>,
 43{
 44    /// Extract Fhe Uint8 ciphertext at `index`
 45    ///
 46    /// `Self` stores batch of Fhe uint8 ciphertext as vector of RLWE
 47    /// ciphertexts. Since Fhe uint8 ciphertext is collection of 8 bool
 48    /// ciphertexts, Fhe uint8 ciphertext at index `i` is stored in coefficients
 49    /// `i*8...(i+1)*8`. To extract Fhe uint8 at index `i`, sample extract bool
 50    /// ciphertext at indices `[i*8, ..., (i+1)*8)`
 51    fn extract_at(&self, index: usize) -> FheUint8<R> {
 52        assert!(index < self.count);
 53        BoolEvaluator::with_local(|e| {
 54            let ring_size = e.parameters().rlwe_n().0;
 55
 56            let start_index = index * 8;
 57            let end_index = (index + 1) * 8;
 58            let data = (start_index..end_index)
 59                .map(|i| {
 60                    let rlwe_index = i / ring_size;
 61                    let coeff_index = i % ring_size;
 62                    self.data[rlwe_index].extract_at(coeff_index)
 63                })
 64                .collect_vec();
 65            FheUint8 { data }
 66        })
 67    }
 68
 69    /// Extracts all FheUint8s packed in vector of RLWE ciphertexts of `Self`
 70    fn extract_all(&self) -> Vec<FheUint8<R>> {
 71        (0..self.count)
 72            .map(|index| self.extract_at(index))
 73            .collect_vec()
 74    }
 75
 76    /// Extracts first `how_many` FheUint8s packed in vector of RLWE
 77    /// ciphertexts of `Self`
 78    fn extract_many(&self, how_many: usize) -> Vec<FheUint8<R>> {
 79        (0..how_many)
 80            .map(|index| self.extract_at(index))
 81            .collect_vec()
 82    }
 83}
 84
 85impl<M: MatrixEntity + MatrixMut<MatElement = u64>> From<&SeededBatchedFheUint8<M::R, [u8; 32]>>
 86    for BatchedFheUint8<M>
 87where
 88    <M as Matrix>::R: RowMut,
 89{
 90    /// Unseeds collection of seeded RLWE ciphertext in SeededBatchedFheUint8
 91    /// and returns as `Self`
 92    fn from(value: &SeededBatchedFheUint8<M::R, [u8; 32]>) -> Self {
 93        BoolEvaluator::with_local(|e| {
 94            let parameters = e.parameters();
 95            let ring_size = parameters.rlwe_n().0;
 96            let rlwe_q = parameters.rlwe_q();
 97
 98            let mut prng = DefaultSecureRng::new_seeded(value.seed);
 99            let rlwes = value
100                .data
101                .iter()
102                .map(|partb| {
103                    let mut rlwe = M::zeros(2, ring_size);
104
105                    // sample A
106                    RandomFillUniformInModulus::random_fill(&mut prng, rlwe_q, rlwe.get_row_mut(0));
107
108                    // Copy over B
109                    rlwe.get_row_mut(1).copy_from_slice(partb.as_ref());
110
111                    rlwe
112                })
113                .collect_vec();
114            Self {
115                data: rlwes,
116                count: value.count,
117            }
118        })
119    }
120}
121
122pub struct SeededBatchedFheUint8<C, S> {
123    /// Vector of Seeded RLWE ciphertexts `C`.
124    ///
125    /// If RLWE(m) = [a, b] s.t. m + e = b - as, `a` can be seeded and seeded
126    /// RLWE ciphertext only contains `b` polynomial
127    data: Vec<C>,
128    /// Seed for the ciphertexts
129    seed: S,
130    /// Count of FheUint8s packed in vector of RLWE ciphertexts
131    count: usize,
132}
133
134impl<C, S> SeededBatchedFheUint8<C, S> {
135    pub fn unseed<M>(&self) -> BatchedFheUint8<M>
136    where
137        BatchedFheUint8<M>: for<'a> From<&'a SeededBatchedFheUint8<C, S>>,
138        M: Matrix<R = C>,
139    {
140        BatchedFheUint8::from(self)
141    }
142}
143
144impl<K, C, S> Encryptor<[u8], SeededBatchedFheUint8<C, S>> for K
145where
146    K: Encryptor<[bool], (Vec<C>, S)>,
147{
148    /// Encrypt a slice of u8s of arbitray length as `SeededBatchedFheUint8`
149    fn encrypt(&self, m: &[u8]) -> SeededBatchedFheUint8<C, S> {
150        // convert vector of u8s to vector bools
151        let bool_m = m
152            .iter()
153            .flat_map(|v| (0..8).into_iter().map(|i| (((*v) >> i) & 1) == 1))
154            .collect_vec();
155        let (cts, seed) = K::encrypt(&self, &bool_m);
156        SeededBatchedFheUint8 {
157            data: cts,
158            seed,
159            count: m.len(),
160        }
161    }
162}
163
164impl<K, C> Encryptor<[u8], BatchedFheUint8<C>> for K
165where
166    K: Encryptor<[bool], Vec<C>>,
167{
168    fn encrypt(&self, m: &[u8]) -> BatchedFheUint8<C> {
169        let bool_m = m
170            .iter()
171            .flat_map(|v| {
172                (0..8)
173                    .into_iter()
174                    .map(|i| ((*v >> i) & 1) == 1)
175                    .collect_vec()
176            })
177            .collect_vec();
178        let cts = K::encrypt(&self, &bool_m);
179        BatchedFheUint8 {
180            data: cts,
181            count: m.len(),
182        }
183    }
184}
185
186impl<C> KeySwitchWithId<BatchedFheUint8<C>> for BatchedFheUint8<C>
187where
188    C: KeySwitchWithId<C>,
189{
190    /// Key switching collection of RLWE ciphertexts in `BatchedFheUint8` from
191    /// user j's RLWE secret u_j to ideal RLWE secret key `s` of the protocol.
192    ///
193    /// - user_id: user id of user j
194    fn key_switch(&self, user_id: usize) -> BatchedFheUint8<C> {
195        let data = self
196            .data
197            .iter()
198            .map(|c| c.key_switch(user_id))
199            .collect_vec();
200        BatchedFheUint8 {
201            data,
202            count: self.count,
203        }
204    }
205}
206
207impl<C, K> MultiPartyDecryptor<u8, FheUint8<C>> for K
208where
209    K: MultiPartyDecryptor<bool, C>,
210    <Self as MultiPartyDecryptor<bool, C>>::DecryptionShare: Clone,
211{
212    type DecryptionShare = Vec<<Self as MultiPartyDecryptor<bool, C>>::DecryptionShare>;
213    fn gen_decryption_share(&self, c: &FheUint8<C>) -> Self::DecryptionShare {
214        assert!(c.data().len() == 8);
215        c.data()
216            .iter()
217            .map(|bit_c| {
218                let decryption_share =
219                    MultiPartyDecryptor::<bool, C>::gen_decryption_share(self, bit_c);
220                decryption_share
221            })
222            .collect_vec()
223    }
224
225    fn aggregate_decryption_shares(&self, c: &FheUint8<C>, shares: &[Self::DecryptionShare]) -> u8 {
226        let mut out = 0u8;
227
228        (0..8).into_iter().for_each(|i| {
229            // Collect bit i^th decryption share of each party
230            let bit_i_decryption_shares = shares.iter().map(|s| s[i].clone()).collect_vec();
231            let bit_i = MultiPartyDecryptor::<bool, C>::aggregate_decryption_shares(
232                self,
233                &c.data()[i],
234                &bit_i_decryption_shares,
235            );
236
237            if bit_i {
238                out += 1 << i;
239            }
240        });
241
242        out
243    }
244}
245
246impl<C, K> Encryptor<u8, FheUint8<C>> for K
247where
248    K: Encryptor<bool, C>,
249{
250    fn encrypt(&self, m: &u8) -> FheUint8<C> {
251        let cts = (0..8)
252            .into_iter()
253            .map(|i| {
254                let bit = ((m >> i) & 1) == 1;
255                K::encrypt(self, &bit)
256            })
257            .collect_vec();
258        FheUint8 { data: cts }
259    }
260}
261
262impl<K, C> Decryptor<u8, FheUint8<C>> for K
263where
264    K: Decryptor<bool, C>,
265{
266    fn decrypt(&self, c: &FheUint8<C>) -> u8 {
267        assert!(c.data.len() == 8);
268        let mut out = 0u8;
269        c.data().iter().enumerate().for_each(|(index, bit_c)| {
270            let bool = K::decrypt(self, bit_c);
271            if bool {
272                out += 1 << index;
273            }
274        });
275        out
276    }
277}