1use itertools::Itertools;
  2use poulpy_core::layouts::prepared::GGSWCiphertextPreparedToRef;
  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 GGSWCiphertextPreparedToRef<BE>> = a
 64        .blocks
 65        .iter()
 66        .map(|x| x as &dyn GGSWCiphertextPreparedToRef<BE>)
 67        .chain(
 68            b.blocks
 69                .iter()
 70                .map(|x| x as &dyn GGSWCiphertextPreparedToRef<BE>),
 71        )
 72        .collect_vec();
 73
 74    // Evaluates out[i] = circuit[i](a, b)
 75    circuit.execute(module, &mut out.blocks, &inputs, scratch);
 76}
 77
 78#[macro_export]
 79macro_rules! define_bdd_2w_to_1w_trait {
 80    ($(#[$meta:meta])* $vis:vis $trait_name:ident, $method_name:ident) => {
 81        $(#[$meta])*
 82        $vis trait $trait_name<T: UnsignedInteger, BE: Backend> {
 83            fn $method_name<A, B>(
 84                &mut self,
 85                module: &Module<BE>,
 86                a: &FheUintBlocksPrep<A, BE, T>,
 87                b: &FheUintBlocksPrep<B, BE, T>,
 88                scratch: &mut Scratch<BE>,
 89            ) where
 90                A: DataRef,
 91                B: DataRef;
 92        }
 93    };
 94}
 95
 96#[macro_export]
 97macro_rules! impl_bdd_2w_to_1w_trait {
 98    ($trait_name:ident, $method_name:ident, $ty:ty, $n:literal, $circuit_ty:ty, $output_circuits:path) => {
 99        impl<D: DataMut, BE: Backend> $trait_name<$ty, BE> for FheUintBlocks<D, $ty>
100        where
101            Circuits2WTo1W<$circuit_ty, $n>: EvalBDD2WTo1W<BE, $ty>,
102        {
103            fn $method_name<A, B>(
104                &mut self,
105                module: &Module<BE>,
106                a: &FheUintBlocksPrep<A, BE, $ty>,
107                b: &FheUintBlocksPrep<B, BE, $ty>,
108                scratch: &mut Scratch<BE>,
109            ) where
110                A: DataRef,
111                B: DataRef,
112            {
113                const OP: Circuits2WTo1W<$circuit_ty, $n> = Circuits2WTo1W::<$circuit_ty, $n>(&$output_circuits);
114
115                OP.eval_bdd_2w_to_1w(module, self, a, b, scratch);
116            }
117        }
118    };
119}
120
121define_bdd_2w_to_1w_trait!(pub Add, add);
122define_bdd_2w_to_1w_trait!(pub Sub, sub);
123define_bdd_2w_to_1w_trait!(pub Sll, sll);
124define_bdd_2w_to_1w_trait!(pub Sra, sra);
125define_bdd_2w_to_1w_trait!(pub Srl, srl);
126define_bdd_2w_to_1w_trait!(pub Slt, slt);
127define_bdd_2w_to_1w_trait!(pub Sltu, sltu);
128define_bdd_2w_to_1w_trait!(pub Or, or);
129define_bdd_2w_to_1w_trait!(pub And, and);
130define_bdd_2w_to_1w_trait!(pub Xor, xor);
131
132impl_bdd_2w_to_1w_trait!(
133    Add,
134    add,
135    u32,
136    32,
137    circuits::u32::add_codegen::AnyBitCircuit,
138    circuits::u32::add_codegen::OUTPUT_CIRCUITS
139);
140
141impl_bdd_2w_to_1w_trait!(
142    Sub,
143    sub,
144    u32,
145    32,
146    circuits::u32::sub_codegen::AnyBitCircuit,
147    circuits::u32::sub_codegen::OUTPUT_CIRCUITS
148);
149
150impl_bdd_2w_to_1w_trait!(
151    Sll,
152    sll,
153    u32,
154    32,
155    circuits::u32::sll_codegen::AnyBitCircuit,
156    circuits::u32::sll_codegen::OUTPUT_CIRCUITS
157);
158
159impl_bdd_2w_to_1w_trait!(
160    Sra,
161    sra,
162    u32,
163    32,
164    circuits::u32::sra_codegen::AnyBitCircuit,
165    circuits::u32::sra_codegen::OUTPUT_CIRCUITS
166);
167
168impl_bdd_2w_to_1w_trait!(
169    Srl,
170    srl,
171    u32,
172    32,
173    circuits::u32::srl_codegen::AnyBitCircuit,
174    circuits::u32::srl_codegen::OUTPUT_CIRCUITS
175);
176
177impl_bdd_2w_to_1w_trait!(
178    Slt,
179    slt,
180    u32,
181    1,
182    circuits::u32::slt_codegen::AnyBitCircuit,
183    circuits::u32::slt_codegen::OUTPUT_CIRCUITS
184);
185
186impl_bdd_2w_to_1w_trait!(
187    Sltu,
188    sltu,
189    u32,
190    1,
191    circuits::u32::sltu_codegen::AnyBitCircuit,
192    circuits::u32::sltu_codegen::OUTPUT_CIRCUITS
193);
194
195impl_bdd_2w_to_1w_trait!(
196    And,
197    and,
198    u32,
199    32,
200    circuits::u32::and_codegen::AnyBitCircuit,
201    circuits::u32::and_codegen::OUTPUT_CIRCUITS
202);
203
204impl_bdd_2w_to_1w_trait!(
205    Or,
206    or,
207    u32,
208    32,
209    circuits::u32::or_codegen::AnyBitCircuit,
210    circuits::u32::or_codegen::OUTPUT_CIRCUITS
211);
212
213impl_bdd_2w_to_1w_trait!(
214    Xor,
215    xor,
216    u32,
217    32,
218    circuits::u32::xor_codegen::AnyBitCircuit,
219    circuits::u32::xor_codegen::OUTPUT_CIRCUITS
220);