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