1use poulpy_core::{
  2    GLWECopy, GLWEPacking, ScratchTakeCore,
  3    layouts::{GGLWEInfos, GGLWEPreparedToRef, GLWEAutomorphismKeyHelper, GetGaloisElement},
  4};
  5use poulpy_hal::{
  6    api::ModuleLogN,
  7    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  8};
  9
 10use crate::tfhe::bdd_arithmetic::{ExecuteBDDCircuit, FheUint, FheUintPrepared, GetBitCircuitInfo, UnsignedInteger, circuits};
 11
 12impl<BE: Backend> ExecuteBDDCircuit1WTo1W<BE> for Module<BE> where Self: Sized + ExecuteBDDCircuit<BE> + GLWEPacking<BE> + GLWECopy
 13{}
 14
 15pub trait ExecuteBDDCircuit1WTo1W<BE: Backend>
 16where
 17    Self: Sized + ModuleLogN + ExecuteBDDCircuit<BE> + GLWEPacking<BE> + GLWECopy,
 18{
 19    fn execute_bdd_circuit_1w_to_1w<R, C, A, K, H, T>(
 20        &self,
 21        out: &mut FheUint<R, T>,
 22        circuit: &C,
 23        a: &FheUintPrepared<A, T, BE>,
 24        key: &H,
 25        scratch: &mut Scratch<BE>,
 26    ) where
 27        T: UnsignedInteger,
 28        C: GetBitCircuitInfo,
 29        R: DataMut,
 30        A: DataRef,
 31        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
 32        H: GLWEAutomorphismKeyHelper<K, BE>,
 33        Scratch<BE>: ScratchTakeCore<BE>,
 34    {
 35        self.execute_bdd_circuit_1w_to_1w_multi_thread(1, out, circuit, a, key, scratch);
 36    }
 37
 38    #[allow(clippy::too_many_arguments)]
 39    /// Operations Z x Z -> Z
 40    fn execute_bdd_circuit_1w_to_1w_multi_thread<R, C, A, K, H, T>(
 41        &self,
 42        threads: usize,
 43        out: &mut FheUint<R, T>,
 44        circuit: &C,
 45        a: &FheUintPrepared<A, T, BE>,
 46        key: &H,
 47        scratch: &mut Scratch<BE>,
 48    ) where
 49        T: UnsignedInteger,
 50        C: GetBitCircuitInfo,
 51        R: DataMut,
 52        A: DataRef,
 53        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
 54        H: GLWEAutomorphismKeyHelper<K, BE>,
 55        Scratch<BE>: ScratchTakeCore<BE>,
 56    {
 57        let (mut out_bits, scratch_1) = scratch.take_glwe_slice(T::BITS as usize, out);
 58
 59        // Evaluates out[i] = circuit[i](a, b)
 60        self.execute_bdd_circuit_multi_thread(threads, &mut out_bits, a, circuit, scratch_1);
 61
 62        // Repacks the bits
 63        out.pack(self, out_bits, key, scratch_1);
 64    }
 65}
 66
 67#[macro_export]
 68macro_rules! define_bdd_1w_to_1w_trait {
 69    ($(#[$meta:meta])* $vis:vis $trait_name:ident, $method_name:ident) => {
 70        paste::paste! {
 71            $(#[$meta])*
 72            $vis trait $trait_name<T: UnsignedInteger, BE: Backend> {
 73
 74                /// Single-threaded version
 75                fn $method_name<A, M, K, H>(
 76                    &mut self,
 77                    module: &M,
 78                    a: &FheUintPrepared<A, T, BE>,
 79                    key: &H,
 80                    scratch: &mut Scratch<BE>,
 81                ) where
 82                    M: ExecuteBDDCircuit1WTo1W<BE>,
 83                    A: DataRef,
 84                    K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
 85                    H: GLWEAutomorphismKeyHelper<K, BE>,
 86                    Scratch<BE>: ScratchTakeCore<BE>;
 87
 88                /// Multithreaded version – same vis, method_name + "_multi_thread"
 89                fn [<$method_name _multi_thread>]<A, M, K, H>(
 90                    &mut self,
 91                    threads: usize,
 92                    module: &M,
 93                    a: &FheUintPrepared<A, T, BE>,
 94                    key: &H,
 95                    scratch: &mut Scratch<BE>,
 96                ) where
 97                    M: ExecuteBDDCircuit1WTo1W<BE>,
 98                    A: DataRef,
 99                    K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
100                    H: GLWEAutomorphismKeyHelper<K, BE>,
101                    Scratch<BE>: ScratchTakeCore<BE>;
102            }
103        }
104    };
105}
106
107#[macro_export]
108macro_rules! impl_bdd_1w_to_1w_trait {
109    ($trait_name:ident, $method_name:ident, $ty:ty, $circuit_ty:ty, $output_circuits:path) => {
110        paste::paste! {
111            impl<D: DataMut, BE: Backend> $trait_name<$ty, BE> for FheUint<D, $ty> {
112
113                fn $method_name<A, M, K, H>(
114                    &mut self,
115                    module: &M,
116                    a: &FheUintPrepared<A, $ty, BE>,
117                    key: &H,
118                    scratch: &mut Scratch<BE>,
119                ) where
120                    M: ExecuteBDDCircuit1WTo1W<BE>,
121                    A: DataRef,
122                    K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
123                    H: GLWEAutomorphismKeyHelper<K, BE>,
124                    Scratch<BE>: ScratchTakeCore<BE>,
125                {
126                    module.execute_bdd_circuit_1w_to_1w(self, &$output_circuits, a, key, scratch)
127                }
128
129                fn [<$method_name _multi_thread>]<A, M, K, H>(
130                    &mut self,
131                    threads: usize,
132                    module: &M,
133                    a: &FheUintPrepared<A, $ty, BE>,
134                    key: &H,
135                    scratch: &mut Scratch<BE>,
136                ) where
137                    M: ExecuteBDDCircuit1WTo1W<BE>,
138                    A: DataRef,
139                    K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
140                    H: GLWEAutomorphismKeyHelper<K, BE>,
141                    Scratch<BE>: ScratchTakeCore<BE>,
142                {
143                    module.execute_bdd_circuit_1w_to_1w_multi_thread(threads, self, &$output_circuits, a, key, scratch)
144                }
145            }
146        }
147    };
148}
149define_bdd_1w_to_1w_trait!(pub Identity, identity);
150
151impl_bdd_1w_to_1w_trait!(
152    Identity,
153    identity,
154    u32,
155    circuits::u32::identity_codgen::AnyBitCircuit,
156    circuits::u32::identity_codgen::OUTPUT_CIRCUITS
157);