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            assert!(nodes.len().is_multiple_of(max_inter_state));
 82
 83            let (mut level, scratch_1) = scratch.take_glwe_slice(max_inter_state * 2, out_i);
 84
 85            level.iter_mut().for_each(|ct| ct.data_mut().zero());
 86
 87            // TODO: implement API on GLWE
 88            level[1]
 89                .data_mut()
 90                .encode_coeff_i64(out_i.base2k().into(), 0, 2, 0, 1);
 91
 92            let mut level_ref = level.iter_mut().collect_vec();
 93            let (mut prev_level, mut next_level) = level_ref.split_at_mut(max_inter_state);
 94
 95            let (all_but_last, last) = nodes.split_at(nodes.len() - max_inter_state);
 96
 97            for nodes_lvl in all_but_last.chunks_exact(max_inter_state) {
 98                for (j, node) in nodes_lvl.iter().enumerate() {
 99                    match node {
100                        Node::Cmux(in_idx, hi_idx, lo_idx) => {
101                            self.cmux(
102                                next_level[j],
103                                prev_level[*hi_idx],
104                                prev_level[*lo_idx],
105                                &inputs.get_bit(*in_idx),
106                                scratch_1,
107                            );
108                        }
109                        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 */
110                        Node::None => {}
111                    }
112                }
113
114                (prev_level, next_level) = (next_level, prev_level);
115            }
116
117            // Last chunck of max_inter_state Nodes is always structured as
118            // [CMUX, NONE, NONE, ..., NONE]
119            match &last[0] {
120                Node::Cmux(in_idx, hi_idx, lo_idx) => {
121                    self.cmux(
122                        out_i,
123                        prev_level[*hi_idx],
124                        prev_level[*lo_idx],
125                        &inputs.get_bit(*in_idx),
126                        scratch_1,
127                    );
128                }
129                _ => {
130                    panic!("invalid last node, should be CMUX")
131                }
132            }
133        }
134
135        for out_i in out.iter_mut().skip(circuit.output_size()) {
136            out_i.data_mut().zero();
137        }
138    }
139}
140
141impl<const N: usize> BitCircuit<N> {
142    pub const fn new(nodes: [Node; N], max_inter_state: usize) -> Self {
143        Self {
144            nodes,
145            max_inter_state,
146        }
147    }
148}
149impl<const N: usize> BitCircuitInfo for BitCircuit<N> {
150    fn info(&self) -> (&[Node], usize) {
151        (self.nodes.as_ref(), self.max_inter_state)
152    }
153}
154
155#[derive(Debug)]
156pub enum Node {
157    Cmux(usize, usize, usize),
158    Copy,
159    None,
160}
161
162pub trait Cmux<BE: Backend>
163where
164    Self: GLWEExternalProduct<BE> + GLWESub + GLWEAdd,
165{
166    fn cmux_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
167    where
168        R: GLWEInfos,
169        A: GLWEInfos,
170        B: GGSWInfos,
171    {
172        self.glwe_external_product_tmp_bytes(res_infos, a_infos, b_infos)
173    }
174
175    fn cmux<R, T, F, S>(&self, res: &mut R, t: &T, f: &F, s: &S, scratch: &mut Scratch<BE>)
176    where
177        R: GLWEToMut,
178        T: GLWEToRef,
179        F: GLWEToRef,
180        S: GGSWPreparedToRef<BE>,
181        Scratch<BE>: ScratchTakeCore<BE>,
182    {
183        self.glwe_sub(res, t, f);
184        self.glwe_external_product_inplace(res, s, scratch);
185        self.glwe_add_inplace(res, f);
186    }
187
188    fn cmux_inplace<R, A, S>(&self, res: &mut R, a: &A, s: &S, scratch: &mut Scratch<BE>)
189    where
190        R: GLWEToMut,
191        A: GLWEToRef,
192        S: GGSWPreparedToRef<BE>,
193        Scratch<BE>: ScratchTakeCore<BE>,
194    {
195        self.glwe_sub_inplace(res, a);
196        self.glwe_external_product_inplace(res, s, scratch);
197        self.glwe_add_inplace(res, a);
198    }
199}
200
201impl<BE: Backend> Cmux<BE> for Module<BE>
202where
203    Self: GLWEExternalProduct<BE> + GLWESub + GLWEAdd,
204    Scratch<BE>: ScratchTakeCore<BE>,
205{
206}