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