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