1use itertools::Itertools;
  2use poulpy_core::layouts::prepared::GGSWPreparedToRef;
  3use poulpy_hal::layouts::{Backend, DataMut, DataRef, Module, Scratch};
  4
  5use crate::tfhe::bdd_arithmetic::{
  6    ExecuteBDDCircuit, FheUintBlocks, FheUintBlocksPrepared, GetBitCircuitInfo, UnsignedInteger, circuits,
  7};
  8
  9impl<T: UnsignedInteger, BE: Backend> ExecuteBDDCircuit2WTo1W<T, BE> for Module<BE> where Self: Sized + ExecuteBDDCircuit<T, BE> {}
 10
 11pub trait ExecuteBDDCircuit2WTo1W<T: UnsignedInteger, BE: Backend>
 12where
 13    Self: Sized + ExecuteBDDCircuit<T, BE>,
 14{
 15    /// Operations Z x Z -> Z
 16    fn execute_bdd_circuit_2w_to_1w<R, C, A, B>(
 17        &self,
 18        out: &mut FheUintBlocks<R, T>,
 19        circuit: &C,
 20        a: &FheUintBlocksPrepared<A, T, BE>,
 21        b: &FheUintBlocksPrepared<B, T, BE>,
 22        scratch: &mut Scratch<BE>,
 23    ) where
 24        C: GetBitCircuitInfo<T>,
 25        R: DataMut,
 26        A: DataRef,
 27        B: DataRef,
 28    {
 29        assert_eq!(out.blocks.len(), T::WORD_SIZE);
 30        assert_eq!(b.blocks.len(), T::WORD_SIZE);
 31        assert_eq!(b.blocks.len(), T::WORD_SIZE);
 32
 33        // Collects inputs into a single array
 34        let inputs: Vec<&dyn GGSWPreparedToRef<BE>> = a
 35            .blocks
 36            .iter()
 37            .map(|x| x as &dyn GGSWPreparedToRef<BE>)
 38            .chain(b.blocks.iter().map(|x| x as &dyn GGSWPreparedToRef<BE>))
 39            .collect_vec();
 40
 41        // Evaluates out[i] = circuit[i](a, b)
 42        self.execute_bdd_circuit(&mut out.blocks, &inputs, circuit, scratch);
 43    }
 44}
 45
 46#[macro_export]
 47macro_rules! define_bdd_2w_to_1w_trait {
 48    ($(#[$meta:meta])* $vis:vis $trait_name:ident, $method_name:ident) => {
 49        $(#[$meta])*
 50        $vis trait $trait_name<T: UnsignedInteger, BE: Backend> {
 51            fn $method_name<A, M, B>(
 52                &mut self,
 53                module: &M,
 54                a: &FheUintBlocksPrepared<A, T, BE>,
 55                b: &FheUintBlocksPrepared<B, T, BE>,
 56                scratch: &mut Scratch<BE>,
 57            ) where
 58                M: ExecuteBDDCircuit2WTo1W<T, BE>,
 59                A: DataRef,
 60                B: DataRef;
 61        }
 62    };
 63}
 64
 65#[macro_export]
 66macro_rules! impl_bdd_2w_to_1w_trait {
 67    ($trait_name:ident, $method_name:ident, $ty:ty, $n:literal, $circuit_ty:ty, $output_circuits:path) => {
 68        impl<D: DataMut, BE: Backend> $trait_name<$ty, BE> for FheUintBlocks<D, $ty> {
 69            fn $method_name<A, M, B>(
 70                &mut self,
 71                module: &M,
 72                a: &FheUintBlocksPrepared<A, $ty, BE>,
 73                b: &FheUintBlocksPrepared<B, $ty, BE>,
 74                scratch: &mut Scratch<BE>,
 75            ) where
 76                M: ExecuteBDDCircuit2WTo1W<$ty, BE>,
 77                A: DataRef,
 78                B: DataRef,
 79            {
 80                module.execute_bdd_circuit_2w_to_1w(self, &$output_circuits, a, b, scratch)
 81            }
 82        }
 83    };
 84}
 85
 86define_bdd_2w_to_1w_trait!(pub Add, add);
 87define_bdd_2w_to_1w_trait!(pub Sub, sub);
 88define_bdd_2w_to_1w_trait!(pub Sll, sll);
 89define_bdd_2w_to_1w_trait!(pub Sra, sra);
 90define_bdd_2w_to_1w_trait!(pub Srl, srl);
 91define_bdd_2w_to_1w_trait!(pub Slt, slt);
 92define_bdd_2w_to_1w_trait!(pub Sltu, sltu);
 93define_bdd_2w_to_1w_trait!(pub Or, or);
 94define_bdd_2w_to_1w_trait!(pub And, and);
 95define_bdd_2w_to_1w_trait!(pub Xor, xor);
 96
 97impl_bdd_2w_to_1w_trait!(
 98    Add,
 99    add,
100    u32,
101    32,
102    circuits::u32::add_codegen::AnyBitCircuit,
103    circuits::u32::add_codegen::OUTPUT_CIRCUITS
104);
105
106impl_bdd_2w_to_1w_trait!(
107    Sub,
108    sub,
109    u32,
110    32,
111    circuits::u32::sub_codegen::AnyBitCircuit,
112    circuits::u32::sub_codegen::OUTPUT_CIRCUITS
113);
114
115impl_bdd_2w_to_1w_trait!(
116    Sll,
117    sll,
118    u32,
119    32,
120    circuits::u32::sll_codegen::AnyBitCircuit,
121    circuits::u32::sll_codegen::OUTPUT_CIRCUITS
122);
123
124impl_bdd_2w_to_1w_trait!(
125    Sra,
126    sra,
127    u32,
128    32,
129    circuits::u32::sra_codegen::AnyBitCircuit,
130    circuits::u32::sra_codegen::OUTPUT_CIRCUITS
131);
132
133impl_bdd_2w_to_1w_trait!(
134    Srl,
135    srl,
136    u32,
137    32,
138    circuits::u32::srl_codegen::AnyBitCircuit,
139    circuits::u32::srl_codegen::OUTPUT_CIRCUITS
140);
141
142impl_bdd_2w_to_1w_trait!(
143    Slt,
144    slt,
145    u32,
146    1,
147    circuits::u32::slt_codegen::AnyBitCircuit,
148    circuits::u32::slt_codegen::OUTPUT_CIRCUITS
149);
150
151impl_bdd_2w_to_1w_trait!(
152    Sltu,
153    sltu,
154    u32,
155    1,
156    circuits::u32::sltu_codegen::AnyBitCircuit,
157    circuits::u32::sltu_codegen::OUTPUT_CIRCUITS
158);
159
160impl_bdd_2w_to_1w_trait!(
161    And,
162    and,
163    u32,
164    32,
165    circuits::u32::and_codegen::AnyBitCircuit,
166    circuits::u32::and_codegen::OUTPUT_CIRCUITS
167);
168
169impl_bdd_2w_to_1w_trait!(
170    Or,
171    or,
172    u32,
173    32,
174    circuits::u32::or_codegen::AnyBitCircuit,
175    circuits::u32::or_codegen::OUTPUT_CIRCUITS
176);
177
178impl_bdd_2w_to_1w_trait!(
179    Xor,
180    xor,
181    u32,
182    32,
183    circuits::u32::xor_codegen::AnyBitCircuit,
184    circuits::u32::xor_codegen::OUTPUT_CIRCUITS
185);