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