1use std::collections::HashMap;
 2
 3use poulpy_core::{
 4    GLWECopy, GLWEDecrypt, ScratchTakeCore,
 5    layouts::{GGSWInfos, GGSWPrepared, GLWE, GLWEInfos, GLWEToMut, GLWEToRef},
 6};
 7use poulpy_hal::layouts::{Backend, Module, Scratch, ZnxZero};
 8
 9use crate::tfhe::bdd_arithmetic::{Cmux, GetGGSWBit, UnsignedInteger};
10
11impl<T: UnsignedInteger, BE: Backend> GLWEBlinSelection<T, BE> for Module<BE> where Self: GLWECopy + Cmux<BE> + GLWEDecrypt<BE> {}
12
13pub trait GLWEBlinSelection<T: UnsignedInteger, BE: Backend>
14where
15    Self: GLWECopy + Cmux<BE> + GLWEDecrypt<BE>,
16{
17    fn glwe_blind_selection_tmp_bytes<R, K>(&self, res_infos: &R, k_infos: &K) -> usize
18    where
19        R: GLWEInfos,
20        K: GGSWInfos,
21    {
22        self.cmux_tmp_bytes(res_infos, res_infos, k_infos) + GLWE::bytes_of_from_infos(res_infos)
23    }
24
25    #[allow(clippy::too_many_arguments)]
26    fn glwe_blind_selection<R, A, K>(
27        &self,
28        res: &mut R,
29        mut a: HashMap<usize, &mut A>,
30        fhe_uint: &K,
31        bit_rsh: usize,
32        bit_mask: usize,
33        scratch: &mut Scratch<BE>,
34    ) where
35        R: GLWEToMut,
36        A: GLWEToMut + GLWEToRef,
37        K: GetGGSWBit<BE>,
38        Scratch<BE>: ScratchTakeCore<BE>,
39    {
40        assert!(bit_rsh + bit_mask <= T::BITS as usize);
41
42        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
43
44        for i in 0..bit_mask {
45            let t: usize = 1 << (bit_mask - i - 1);
46
47            let bit: &GGSWPrepared<&[u8], BE> = &fhe_uint.get_bit(bit_rsh + bit_mask - i - 1); // MSB -> LSB traversal
48
49            for j in 0..t {
50                let hi: Option<&mut A> = a.remove(&j);
51                let lo: Option<&mut A> = a.remove(&(j + t));
52
53                match (lo, hi) {
54                    (Some(lo), Some(hi)) => {
55                        self.cmux_inplace(lo, hi, bit, scratch);
56                        a.insert(j, lo);
57                    }
58
59                    (Some(lo), None) => {
60                        let (mut zero, scratch_1) = scratch.take_glwe(res);
61                        zero.data_mut().zero();
62                        self.cmux_inplace(lo, &zero, bit, scratch_1);
63                        a.insert(j, lo);
64                    }
65
66                    (None, Some(hi)) => {
67                        let (mut zero, scratch_1) = scratch.take_glwe(res);
68                        zero.data_mut().zero();
69                        self.cmux_inplace(&mut zero, hi, bit, scratch_1);
70                        self.glwe_copy(hi, &zero);
71                        a.insert(j, hi);
72                    }
73
74                    (None, None) => {
75                        // No low or high branch — nothing to insert
76                        // leave empty; future iterations will combine actual ciphertexts
77                    }
78                }
79            }
80        }
81
82        let out: Option<&mut A> = a.remove(&0);
83
84        if let Some(out) = out {
85            self.glwe_copy(res, out);
86        } else {
87            res.data_mut().zero();
88        }
89    }
90}