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