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