1use itertools::Itertools;
  2
  3use crate::{
  4    bool::{BoolEvaluator, FheBool},
  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    data: Vec<C>,
 35}
 36
 37impl<C, R> SampleExtractor<FheUint8<R>> for BatchedFheUint8<C>
 38where
 39    C: SampleExtractor<R>,
 40{
 41    /// Extract Fhe Uint8 ciphertext at `index`
 42    ///
 43    /// `Self` stores batch of Fhe uint8 ciphertext as vector of RLWE
 44    /// ciphertexts. Since Fhe uint8 ciphertext is collection of 8 bool
 45    /// ciphertexts, Fhe uint8 ciphertext at index `i` is stored in coefficients
 46    /// `i*8...(i+1)*8`. To extract Fhe uint8 at index `i`, sample extract bool
 47    /// ciphertext at indices `[i*8, ..., (i+1)*8)`
 48    fn extract(&self, index: usize) -> FheUint8<R> {
 49        BoolEvaluator::with_local(|e| {
 50            let ring_size = e.parameters().rlwe_n().0;
 51
 52            let start_index = index * 8;
 53            let end_index = (index + 1) * 8;
 54            let data = (start_index..end_index)
 55                .map(|i| {
 56                    let rlwe_index = i / ring_size;
 57                    let coeff_index = i % ring_size;
 58                    self.data[rlwe_index].extract(coeff_index)
 59                })
 60                .collect_vec();
 61            FheUint8 { data }
 62        })
 63    }
 64}
 65
 66impl<M: MatrixEntity + MatrixMut<MatElement = u64>> From<&SeededBatchedFheUint8<M::R, [u8; 32]>>
 67    for BatchedFheUint8<M>
 68where
 69    <M as Matrix>::R: RowMut,
 70{
 71    /// Unseeds collection of seeded RLWE ciphertext in SeededBatchedFheUint8
 72    /// and returns as `Self`
 73    fn from(value: &SeededBatchedFheUint8<M::R, [u8; 32]>) -> Self {
 74        BoolEvaluator::with_local(|e| {
 75            let parameters = e.parameters();
 76            let ring_size = parameters.rlwe_n().0;
 77            let rlwe_q = parameters.rlwe_q();
 78
 79            let mut prng = DefaultSecureRng::new_seeded(value.seed);
 80            let rlwes = value
 81                .data
 82                .iter()
 83                .map(|partb| {
 84                    let mut rlwe = M::zeros(2, ring_size);
 85
 86                    // sample A
 87                    RandomFillUniformInModulus::random_fill(&mut prng, rlwe_q, rlwe.get_row_mut(0));
 88
 89                    // Copy over B
 90                    rlwe.get_row_mut(1).copy_from_slice(partb.as_ref());
 91
 92                    rlwe
 93                })
 94                .collect_vec();
 95            Self { data: rlwes }
 96        })
 97    }
 98}
 99
100pub struct SeededBatchedFheUint8<C, S> {
101    data: Vec<C>,
102    seed: S,
103}
104
105impl<C, S> SeededBatchedFheUint8<C, S> {
106    pub fn unseed<M>(&self) -> BatchedFheUint8<M>
107    where
108        BatchedFheUint8<M>: for<'a> From<&'a SeededBatchedFheUint8<C, S>>,
109        M: Matrix<R = C>,
110    {
111        BatchedFheUint8::from(self)
112    }
113}
114
115impl<K, C, S> Encryptor<[u8], SeededBatchedFheUint8<C, S>> for K
116where
117    K: Encryptor<[bool], (Vec<C>, S)>,
118{
119    /// Encrypt a slice of u8s of arbitray length as `SeededBatchedFheUint8`
120    fn encrypt(&self, m: &[u8]) -> SeededBatchedFheUint8<C, S> {
121        // convert vector of u8s to vector bools
122        let m = m
123            .iter()
124            .flat_map(|v| (0..8).into_iter().map(|i| (((*v) >> i) & 1) == 1))
125            .collect_vec();
126        let (cts, seed) = K::encrypt(&self, &m);
127        SeededBatchedFheUint8 { data: cts, seed }
128    }
129}
130
131impl<K, C> Encryptor<[u8], BatchedFheUint8<C>> for K
132where
133    K: Encryptor<[bool], Vec<C>>,
134{
135    fn encrypt(&self, m: &[u8]) -> BatchedFheUint8<C> {
136        let m = m
137            .iter()
138            .flat_map(|v| {
139                (0..8)
140                    .into_iter()
141                    .map(|i| ((*v >> i) & 1) == 1)
142                    .collect_vec()
143            })
144            .collect_vec();
145        let cts = K::encrypt(&self, &m);
146        BatchedFheUint8 { data: cts }
147    }
148}
149
150impl<C> KeySwitchWithId<BatchedFheUint8<C>> for BatchedFheUint8<C>
151where
152    C: KeySwitchWithId<C>,
153{
154    /// Key switching collection of RLWE ciphertexts in `BatchedFheUint8` from
155    /// user j's RLWE secret u_j to ideal RLWE secret key `s` of the protocol.
156    ///
157    /// - user_id: user id of user j
158    fn key_switch(&self, user_id: usize) -> BatchedFheUint8<C> {
159        let data = self
160            .data
161            .iter()
162            .map(|c| c.key_switch(user_id))
163            .collect_vec();
164        BatchedFheUint8 { data }
165    }
166}
167
168impl<C, K> MultiPartyDecryptor<u8, FheUint8<C>> for K
169where
170    K: MultiPartyDecryptor<bool, C>,
171    <Self as MultiPartyDecryptor<bool, C>>::DecryptionShare: Clone,
172{
173    type DecryptionShare = Vec<<Self as MultiPartyDecryptor<bool, C>>::DecryptionShare>;
174    fn gen_decryption_share(&self, c: &FheUint8<C>) -> Self::DecryptionShare {
175        assert!(c.data().len() == 8);
176        c.data()
177            .iter()
178            .map(|bit_c| {
179                let decryption_share =
180                    MultiPartyDecryptor::<bool, C>::gen_decryption_share(self, bit_c);
181                decryption_share
182            })
183            .collect_vec()
184    }
185
186    fn aggregate_decryption_shares(&self, c: &FheUint8<C>, shares: &[Self::DecryptionShare]) -> u8 {
187        let mut out = 0u8;
188
189        (0..8).into_iter().for_each(|i| {
190            // Collect bit i^th decryption share of each party
191            let bit_i_decryption_shares = shares.iter().map(|s| s[i].clone()).collect_vec();
192            let bit_i = MultiPartyDecryptor::<bool, C>::aggregate_decryption_shares(
193                self,
194                &c.data()[i],
195                &bit_i_decryption_shares,
196            );
197
198            if bit_i {
199                out += 1 << i;
200            }
201        });
202
203        out
204    }
205}
206
207impl<C, K> Encryptor<u8, FheUint8<C>> for K
208where
209    K: Encryptor<bool, C>,
210{
211    fn encrypt(&self, m: &u8) -> FheUint8<C> {
212        let cts = (0..8)
213            .into_iter()
214            .map(|i| {
215                let bit = ((m >> i) & 1) == 1;
216                K::encrypt(self, &bit)
217            })
218            .collect_vec();
219        FheUint8 { data: cts }
220    }
221}
222
223impl<K, C> Decryptor<u8, FheUint8<C>> for K
224where
225    K: Decryptor<bool, C>,
226{
227    fn decrypt(&self, c: &FheUint8<C>) -> u8 {
228        assert!(c.data.len() == 8);
229        let mut out = 0u8;
230        c.data().iter().enumerate().for_each(|(index, bit_c)| {
231            let bool = K::decrypt(self, bit_c);
232            if bool {
233                out += 1 << index;
234            }
235        });
236        out
237    }
238}