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