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