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