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