1use core::panic;
  2
  3use itertools::Itertools;
  4use poulpy_core::{
  5    GLWEAdd, GLWECopy, GLWEExternalProduct, GLWENormalize, GLWESub, ScratchTakeCore, layouts::{GGSWInfos, GLWE, GLWEInfos, GLWEToMut, GLWEToRef, LWEInfos, prepared::GGSWPreparedToRef}
  6};
  7use poulpy_hal::layouts::{Backend, DataMut, Module, Scratch, ZnxZero};
  8
  9use crate::tfhe::bdd_arithmetic::{GetGGSWBit, UnsignedInteger};
 10
 11pub trait BitCircuitInfo {
 12    fn info(&self) -> (&[Node], usize);
 13}
 14
 15pub trait GetBitCircuitInfo<T: UnsignedInteger> {
 16    fn input_size(&self) -> usize;
 17    fn output_size(&self) -> usize;
 18    fn get_circuit(&self, bit: usize) -> (&[Node], usize);
 19}
 20
 21pub struct BitCircuit<const N: usize> {
 22    pub nodes: [Node; N],
 23    pub max_inter_state: usize,
 24}
 25
 26pub trait BitCircuitFamily {
 27    const INPUT_BITS: usize;
 28    const OUTPUT_BITS: usize;
 29}
 30
 31pub struct Circuit<C: BitCircuitInfo, const N: usize>(pub [C; N]);
 32
 33impl<C, T: UnsignedInteger, const N: usize> GetBitCircuitInfo<T> for Circuit<C, N>
 34where
 35    C: BitCircuitInfo + BitCircuitFamily,
 36{
 37    fn input_size(&self) -> usize {
 38        C::INPUT_BITS
 39    }
 40    fn output_size(&self) -> usize {
 41        C::OUTPUT_BITS
 42    }
 43    fn get_circuit(&self, bit: usize) -> (&[Node], usize) {
 44        self.0[bit].info()
 45    }
 46}
 47
 48pub trait ExecuteBDDCircuit<T: UnsignedInteger, BE: Backend> {
 49    fn execute_bdd_circuit<C, G, O>(&self, out: &mut [GLWE<O>], inputs: &G, circuit: &C, scratch: &mut Scratch<BE>)
 50    where
 51        G: GetGGSWBit<BE> + BitSize,
 52        C: GetBitCircuitInfo<T>,
 53        O: DataMut;
 54}
 55
 56pub trait BitSize {
 57    fn bit_size(&self) -> usize;
 58}
 59
 60impl<T: UnsignedInteger, BE: Backend> ExecuteBDDCircuit<T, BE> for Module<BE>
 61where
 62    Self: Cmux<BE> + GLWECopy,
 63    Scratch<BE>: ScratchTakeCore<BE>,
 64{
 65    fn execute_bdd_circuit<C, G, O>(&self, out: &mut [GLWE<O>], inputs: &G, circuit: &C, scratch: &mut Scratch<BE>)
 66    where
 67        G: GetGGSWBit<BE> + BitSize,
 68        C: GetBitCircuitInfo<T>,
 69        O: DataMut,
 70    {
 71        #[cfg(debug_assertions)]
 72        {
 73            assert!(inputs.bit_size() >= circuit.input_size(), "inputs.bit_size(): {} < circuit.input_size():{}", inputs.bit_size(), circuit.input_size());
 74            assert!(out.len() >= circuit.output_size(), "out.len(): {} < circuit.output_size(): {}", out.len(), circuit.output_size());
 75        }
 76
 77        for (i, out_i) in out.iter_mut().enumerate().take(circuit.output_size()) {
 78            let (nodes, max_inter_state) = circuit.get_circuit(i);
 79
 80            if max_inter_state == 0 {
 81                out_i.data_mut().zero();
 82            } else {
 83                assert!(nodes.len().is_multiple_of(max_inter_state));
 84
 85                let (mut level, scratch_1) = scratch.take_glwe_slice(max_inter_state * 2, out_i);
 86
 87                level.iter_mut().for_each(|ct| ct.data_mut().zero());
 88
 89                // TODO: implement API on GLWE
 90                level[1]
 91                    .data_mut()
 92                    .encode_coeff_i64(out_i.base2k().into(), 0, 2, 0, 1);
 93
 94                let mut level_ref = level.iter_mut().collect_vec();
 95                let (mut prev_level, mut next_level) = level_ref.split_at_mut(max_inter_state);
 96
 97                let (all_but_last, last) = nodes.split_at(nodes.len() - max_inter_state);
 98
 99                for nodes_lvl in all_but_last.chunks_exact(max_inter_state) {
100                    for (j, node) in nodes_lvl.iter().enumerate() {
101                        match node {
102                            Node::Cmux(in_idx, hi_idx, lo_idx) => {
103                                self.cmux(
104                                    next_level[j],
105                                    prev_level[*hi_idx],
106                                    prev_level[*lo_idx],
107                                    &inputs.get_bit(*in_idx),
108                                    scratch_1,
109                                );
110                            }
111                            Node::Copy => self.glwe_copy(next_level[j], prev_level[j]), /* Update BDD circuits to order Cmux -> Copy -> None so that mem swap can be used */
112                            Node::None => {}
113                        }
114                    }
115
116                    (prev_level, next_level) = (next_level, prev_level);
117                }
118
119                // Last chunck of max_inter_state Nodes is always structured as
120                // [CMUX, NONE, NONE, ..., NONE]
121                match &last[0] {
122                    Node::Cmux(in_idx, hi_idx, lo_idx) => {
123                        self.cmux(
124                            out_i,
125                            prev_level[*hi_idx],
126                            prev_level[*lo_idx],
127                            &inputs.get_bit(*in_idx),
128                            scratch_1,
129                        );
130                    }
131                    _ => {
132                        panic!("invalid last node, should be CMUX")
133                    }
134                }
135            }
136        }
137
138        for out_i in out.iter_mut().skip(circuit.output_size()) {
139            out_i.data_mut().zero();
140        }
141    }
142}
143
144impl<const N: usize> BitCircuit<N> {
145    pub const fn new(nodes: [Node; N], max_inter_state: usize) -> Self {
146        Self {
147            nodes,
148            max_inter_state,
149        }
150    }
151}
152impl<const N: usize> BitCircuitInfo for BitCircuit<N> {
153    fn info(&self) -> (&[Node], usize) {
154        (self.nodes.as_ref(), self.max_inter_state)
155    }
156}
157
158#[derive(Debug)]
159pub enum Node {
160    Cmux(usize, usize, usize),
161    Copy,
162    None,
163}
164
165pub trait Cmux<BE: Backend>
166where
167    Self: GLWEExternalProduct<BE> + GLWESub + GLWEAdd + GLWENormalize<BE>,
168{
169    fn cmux_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
170    where
171        R: GLWEInfos,
172        A: GLWEInfos,
173        B: GGSWInfos,
174    {
175        self.glwe_external_product_tmp_bytes(res_infos, a_infos, b_infos)
176    }
177
178    fn cmux<R, T, F, S>(&self, res: &mut R, t: &T, f: &F, s: &S, scratch: &mut Scratch<BE>)
179    where
180        R: GLWEToMut,
181        T: GLWEToRef,
182        F: GLWEToRef,
183        S: GGSWPreparedToRef<BE>,
184        Scratch<BE>: ScratchTakeCore<BE>,
185    {
186        self.glwe_sub(res, t, f);
187        self.glwe_normalize_inplace(res, scratch);
188        self.glwe_external_product_inplace(res, s, scratch);
189        self.glwe_add_inplace(res, f);
190        self.glwe_normalize_inplace(res, scratch);
191    }
192
193    fn cmux_inplace<R, A, S>(&self, res: &mut R, a: &A, s: &S, scratch: &mut Scratch<BE>)
194    where
195        R: GLWEToMut,
196        A: GLWEToRef,
197        S: GGSWPreparedToRef<BE>,
198        Scratch<BE>: ScratchTakeCore<BE>,
199    {
200        self.glwe_sub_inplace(res, a);
201        self.glwe_normalize_inplace(res, scratch);
202        self.glwe_external_product_inplace(res, s, scratch);
203        self.glwe_add_inplace(res, a);
204        self.glwe_normalize_inplace(res, scratch);
205    }
206}
207
208impl<BE: Backend> Cmux<BE> for Module<BE>
209where
210    Self: GLWEExternalProduct<BE> + GLWESub + GLWEAdd + GLWENormalize<BE>,
211    Scratch<BE>: ScratchTakeCore<BE>,
212{
213}