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