1use itertools::Itertools;
  2use poulpy_core::{
  3    GLWEAdd, GLWECopy, GLWEExternalProduct, GLWESub, ScratchTakeCore,
  4    layouts::{
  5        GLWE, LWEInfos,
  6        prepared::{GGSWPrepared, GGSWPreparedToRef},
  7    },
  8};
  9use poulpy_hal::layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxZero};
 10
 11use crate::tfhe::bdd_arithmetic::UnsignedInteger;
 12
 13pub trait BitCircuitInfo {
 14    fn info(&self) -> (&[Node], &[usize], usize);
 15}
 16
 17pub trait GetBitCircuitInfo<T: UnsignedInteger> {
 18    fn input_size(&self) -> usize;
 19    fn output_size(&self) -> usize;
 20    fn get_circuit(&self, bit: usize) -> (&[Node], &[usize], usize);
 21}
 22
 23pub(crate) struct BitCircuit<const N: usize, const K: usize> {
 24    pub(crate) nodes: [Node; N],
 25    pub(crate) levels: [usize; K],
 26    pub(crate) max_inter_state: usize,
 27}
 28
 29pub struct Circuit<C: BitCircuitInfo, const N: usize>(pub [C; N]);
 30
 31pub trait ExecuteBDDCircuit<T: UnsignedInteger, BE: Backend> {
 32    fn execute_bdd_circuit<C, O>(
 33        &self,
 34        out: &mut [GLWE<O>],
 35        inputs: &[&dyn GGSWPreparedToRef<BE>],
 36        circuit: &C,
 37        scratch: &mut Scratch<BE>,
 38    ) where
 39        C: GetBitCircuitInfo<T>,
 40        O: DataMut;
 41}
 42
 43impl<T: UnsignedInteger, BE: Backend> ExecuteBDDCircuit<T, BE> for Module<BE>
 44where
 45    Self: Cmux<BE> + GLWECopy,
 46    Scratch<BE>: ScratchTakeCore<BE>,
 47{
 48    fn execute_bdd_circuit<C, O>(
 49        &self,
 50        out: &mut [GLWE<O>],
 51        inputs: &[&dyn GGSWPreparedToRef<BE>],
 52        circuit: &C,
 53        scratch: &mut Scratch<BE>,
 54    ) where
 55        C: GetBitCircuitInfo<T>,
 56        O: DataMut,
 57    {
 58        #[cfg(debug_assertions)]
 59        {
 60            assert_eq!(inputs.len(), circuit.input_size());
 61            assert!(out.len() >= circuit.output_size());
 62        }
 63
 64        for (i, out_i) in out.iter_mut().enumerate().take(circuit.output_size()) {
 65            let (nodes, levels, max_inter_state) = circuit.get_circuit(i);
 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            for i in 0..levels.len() - 1 {
 80                let start: usize = levels[i];
 81                let end: usize = levels[i + 1];
 82
 83                let nodes_lvl: &[Node] = &nodes[start..end];
 84
 85                for (j, node) in nodes_lvl.iter().enumerate() {
 86                    if node.low_index == node.high_index {
 87                        self.glwe_copy(next_level[j], prev_level[node.low_index]);
 88                    } else {
 89                        self.cmux(
 90                            next_level[j],
 91                            prev_level[node.high_index],
 92                            prev_level[node.low_index],
 93                            &inputs[node.input_index].to_ref(),
 94                            scratch_1,
 95                        );
 96                    }
 97                }
 98
 99                (prev_level, next_level) = (next_level, prev_level);
100            }
101
102            // handle last output
103            // there's always only 1 node at last level
104            let node: &Node = nodes.last().unwrap();
105            self.cmux(
106                out_i,
107                prev_level[node.high_index],
108                prev_level[node.low_index],
109                &inputs[node.input_index].to_ref(),
110                scratch_1,
111            );
112        }
113
114        for out_i in out.iter_mut().skip(circuit.output_size()) {
115            out_i.data_mut().zero();
116        }
117    }
118}
119
120impl<const N: usize, const K: usize> BitCircuit<N, K> {
121    pub(crate) const fn new(nodes: [Node; N], levels: [usize; K], max_inter_state: usize) -> Self {
122        Self {
123            nodes,
124            levels,
125            max_inter_state,
126        }
127    }
128}
129impl<const N: usize, const K: usize> BitCircuitInfo for BitCircuit<N, K> {
130    fn info(&self) -> (&[Node], &[usize], usize) {
131        (
132            self.nodes.as_ref(),
133            self.levels.as_ref(),
134            self.max_inter_state,
135        )
136    }
137}
138
139#[derive(Debug)]
140pub struct Node {
141    input_index: usize,
142    high_index: usize,
143    low_index: usize,
144}
145
146impl Node {
147    pub(crate) const fn new(input_index: usize, high_index: usize, low_index: usize) -> Self {
148        Self {
149            input_index,
150            high_index,
151            low_index,
152        }
153    }
154}
155
156pub trait Cmux<BE: Backend> {
157    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>)
158    where
159        O: DataMut,
160        T: DataRef,
161        F: DataRef,
162        S: DataRef;
163}
164
165impl<BE: Backend> Cmux<BE> for Module<BE>
166where
167    Module<BE>: GLWEExternalProduct<BE> + GLWESub + GLWEAdd,
168    Scratch<BE>: ScratchTakeCore<BE>,
169{
170    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>)
171    where
172        O: DataMut,
173        T: DataRef,
174        F: DataRef,
175        S: DataRef,
176    {
177        // let mut out: GLWECiphertext<&mut [u8]> = out.to_mut();
178        self.glwe_sub(out, t, f);
179        self.glwe_external_product_inplace(out, s, scratch);
180        self.glwe_add_inplace(out, f);
181    }
182}