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