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