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