1use itertools::Itertools;
  2use poulpy_core::{
  3    GLWECopy, ScratchTakeCore,
  4    layouts::{GGSWInfos, GGSWPrepared, GLWE, GLWEInfos, GLWEToMut, GLWEToRef},
  5};
  6use poulpy_hal::layouts::{Backend, Module, Scratch};
  7
  8use crate::tfhe::bdd_arithmetic::{Cmux, Cswap, GetGGSWBit};
  9
 10pub struct GLWEBlindRetriever {
 11    accumulators: Vec<Accumulator>,
 12    counter: usize,
 13}
 14
 15impl GLWEBlindRetriever {
 16    pub fn alloc<A>(infos: &A, size: usize) -> Self
 17    where
 18        A: GLWEInfos,
 19    {
 20        let bit_size: usize = (u32::BITS - (size as u32 - 1).leading_zeros()) as usize;
 21        Self {
 22            accumulators: (0..bit_size)
 23                .map(|_| Accumulator::alloc(infos))
 24                .collect_vec(),
 25            counter: 0,
 26        }
 27    }
 28
 29    pub fn retrieve_tmp_bytes<M, R, S, BE: Backend>(module: &M, res: &R, selector: &S) -> usize
 30    where
 31        M: Cmux<BE>,
 32        R: GLWEInfos,
 33        S: GGSWInfos,
 34    {
 35        module.cmux_tmp_bytes(res, res, selector)
 36    }
 37
 38    pub fn retrieve<M, R, A, S, BE: Backend>(
 39        &mut self,
 40        module: &M,
 41        res: &mut R,
 42        data: &[A],
 43        selector: &S,
 44        offset: usize,
 45        scratch: &mut Scratch<BE>,
 46    ) where
 47        M: GLWECopy + Cmux<BE>,
 48        R: GLWEToMut,
 49        A: GLWEToRef,
 50        S: GetGGSWBit<BE>,
 51        Scratch<BE>: ScratchTakeCore<BE>,
 52    {
 53        self.reset();
 54        for ct in data {
 55            self.add(module, ct, selector, offset, scratch);
 56        }
 57        self.flush(module, res, selector, offset, scratch);
 58    }
 59
 60    pub fn add<A, S, M, BE: Backend>(&mut self, module: &M, a: &A, selector: &S, offset: usize, scratch: &mut Scratch<BE>)
 61    where
 62        A: GLWEToRef,
 63        S: GetGGSWBit<BE>,
 64        M: GLWECopy + Cmux<BE>,
 65        Scratch<BE>: ScratchTakeCore<BE>,
 66    {
 67        assert!(
 68            (self.counter as u32) < 1 << self.accumulators.len(),
 69            "Accumulating limit of {} reached",
 70            1 << self.accumulators.len()
 71        );
 72
 73        add_core(
 74            module,
 75            a,
 76            &mut self.accumulators,
 77            0,
 78            selector,
 79            offset,
 80            scratch,
 81        );
 82        self.counter += 1;
 83    }
 84
 85    pub fn flush<R, M, S, BE: Backend>(&mut self, module: &M, res: &mut R, selector: &S, offset: usize, scratch: &mut Scratch<BE>)
 86    where
 87        R: GLWEToMut,
 88        S: GetGGSWBit<BE>,
 89        M: GLWECopy + Cmux<BE>,
 90        Scratch<BE>: ScratchTakeCore<BE>,
 91    {
 92        for i in 0..self.accumulators.len() - 1 {
 93            let (acc_prev, acc_next) = self.accumulators.split_at_mut(i + 1);
 94            if acc_prev[i].num != 0 {
 95                add_core(
 96                    module,
 97                    &acc_prev[i].data,
 98                    acc_next,
 99                    i + 1,
100                    selector,
101                    offset,
102                    scratch,
103                );
104                acc_prev[0].num = 0
105            }
106        }
107        module.glwe_copy(res, &self.accumulators.last().unwrap().data);
108        self.reset()
109    }
110
111    fn reset(&mut self) {
112        for acc in self.accumulators.iter_mut() {
113            acc.num = 0;
114        }
115        self.counter = 0;
116    }
117}
118
119struct Accumulator {
120    data: GLWE<Vec<u8>>,
121    num: usize, // Number of accumulated values
122}
123
124impl Accumulator {
125    pub fn alloc<A>(infos: &A) -> Self
126    where
127        A: GLWEInfos,
128    {
129        Self {
130            data: GLWE::alloc_from_infos(infos),
131            num: 0,
132        }
133    }
134}
135
136fn add_core<A, S, M, BE: Backend>(
137    module: &M,
138    a: &A,
139    accumulators: &mut [Accumulator],
140    i: usize,
141    selector: &S,
142    offset: usize,
143    scratch: &mut Scratch<BE>,
144) where
145    A: GLWEToRef,
146    S: GetGGSWBit<BE>,
147    M: GLWECopy + Cmux<BE>,
148    Scratch<BE>: ScratchTakeCore<BE>,
149{
150    // Isolate the first accumulator
151    let (acc_prev, acc_next) = accumulators.split_at_mut(1);
152
153    match acc_prev[0].num {
154        0 => {
155            module.glwe_copy(&mut acc_prev[0].data, a);
156            acc_prev[0].num = 1;
157        }
158        1 => {
159            module.cmux_inplace_neg(
160                &mut acc_prev[0].data,
161                a,
162                &selector.get_bit(i + offset),
163                scratch,
164            );
165
166            if !acc_next.is_empty() {
167                add_core(
168                    module,
169                    &acc_prev[0].data,
170                    acc_next,
171                    i + 1,
172                    selector,
173                    offset,
174                    scratch,
175                );
176            }
177
178            acc_prev[0].num = 0
179        }
180        _ => {
181            panic!("something went wrong")
182        }
183    }
184}
185
186impl<BE: Backend> GLWEBlindRetrieval<BE> for Module<BE> where Self: GLWECopy + Cmux<BE> + Cswap<BE> {}
187
188pub trait GLWEBlindRetrieval<BE: Backend>
189where
190    Self: GLWECopy + Cmux<BE> + Cswap<BE>,
191{
192    fn glwe_blind_retrieval_tmp_bytes<R, K>(&self, res_infos: &R, k_infos: &K) -> usize
193    where
194        R: GLWEInfos,
195        K: GGSWInfos,
196    {
197        self.cswap_tmp_bytes(res_infos, res_infos, k_infos)
198    }
199
200    fn glwe_blind_retrieval_statefull<R, K>(
201        &self,
202        res: &mut Vec<R>,
203        bits: &K,
204        bit_rsh: usize,
205        bit_mask: usize,
206        scratch: &mut Scratch<BE>,
207    ) where
208        R: GLWEToMut + GLWEInfos,
209        K: GetGGSWBit<BE>,
210        Scratch<BE>: ScratchTakeCore<BE>,
211    {
212        for i in 0..bit_mask {
213            let t: usize = 1 << (bit_mask - i - 1);
214            let bit: &GGSWPrepared<&[u8], BE> = &bits.get_bit(bit_rsh + bit_mask - i - 1); // MSB -> LSB traversal
215            for j in 0..t {
216                if j + t < res.len() {
217                    let (lo, hi) = res.split_at_mut(j + t);
218                    self.cswap(&mut lo[j], &mut hi[0], bit, scratch);
219                }
220            }
221        }
222    }
223
224    fn glwe_blind_retrieval_statefull_rev<R, K>(
225        &self,
226        res: &mut Vec<R>,
227        bits: &K,
228        bit_rsh: usize,
229        bit_mask: usize,
230        scratch: &mut Scratch<BE>,
231    ) where
232        R: GLWEToMut + GLWEInfos,
233        K: GetGGSWBit<BE>,
234        Scratch<BE>: ScratchTakeCore<BE>,
235    {
236        for i in (0..bit_mask).rev() {
237            let t: usize = 1 << (bit_mask - i - 1);
238            let bit: &GGSWPrepared<&[u8], BE> = &bits.get_bit(bit_rsh + bit_mask - i - 1); // MSB -> LSB traversal
239            for j in 0..t {
240                if j < res.len() && j + t < res.len() {
241                    let (lo, hi) = res.split_at_mut(j + t);
242                    self.cswap(&mut lo[j], &mut hi[0], bit, scratch);
243                }
244            }
245        }
246    }
247}