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