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