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::bin_fhe::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
101#[macro_export]
102macro_rules! define_bdd_2w_to_1w_trait {
103    ($(#[$meta:meta])* $vis:vis $trait_name:ident, $method_name:ident) => {
104        paste::paste! {
105            $(#[$meta])*
106            $vis trait $trait_name<T: UnsignedInteger, BE: Backend> {
107
108                /// Single-threaded version
109                fn $method_name<A, M, K, H, B>(
110                    &mut self,
111                    module: &M,
112                    a: &FheUintPrepared<A, T, BE>,
113                    b: &FheUintPrepared<B, T, BE>,
114                    key: &H,
115                    scratch: &mut Scratch<BE>,
116                ) where
117                    M: ExecuteBDDCircuit2WTo1W<BE>,
118                    A: DataRef,
119                    B: DataRef,
120                    K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
121                    H: GLWEAutomorphismKeyHelper<K, BE>,
122                    Scratch<BE>: ScratchTakeCore<BE>;
123
124                /// Multithreaded version – same vis, method_name + "_multi_thread"
125                fn [<$method_name _multi_thread>]<A, M, K, H, B>(
126                    &mut self,
127                    threads: usize,
128                    module: &M,
129                    a: &FheUintPrepared<A, T, BE>,
130                    b: &FheUintPrepared<B, T, BE>,
131                    key: &H,
132                    scratch: &mut Scratch<BE>,
133                ) where
134                    M: ExecuteBDDCircuit2WTo1W<BE>,
135                    A: DataRef,
136                    B: DataRef,
137                    K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
138                    H: GLWEAutomorphismKeyHelper<K, BE>,
139                    Scratch<BE>: ScratchTakeCore<BE>;
140            }
141        }
142    };
143}
144
145#[macro_export]
146macro_rules! impl_bdd_2w_to_1w_trait {
147    ($trait_name:ident, $method_name:ident, $ty:ty, $circuit_ty:ty, $output_circuits:path) => {
148        paste::paste! {
149            impl<D: DataMut, BE: Backend> $trait_name<$ty, BE> for FheUint<D, $ty> {
150
151                fn $method_name<A, M, K, H, B>(
152                    &mut self,
153                    module: &M,
154                    a: &FheUintPrepared<A, $ty, BE>,
155                    b: &FheUintPrepared<B, $ty, BE>,
156                    key: &H,
157                    scratch: &mut Scratch<BE>,
158                ) where
159                    M: ExecuteBDDCircuit2WTo1W<BE>,
160                    A: DataRef,
161                    B: DataRef,
162                    K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
163                    H: GLWEAutomorphismKeyHelper<K, BE>,
164                    Scratch<BE>: ScratchTakeCore<BE>,
165                {
166                    module.execute_bdd_circuit_2w_to_1w(self, &$output_circuits, a, b, key, scratch)
167                }
168
169                fn [<$method_name _multi_thread>]<A, M, K, H, B>(
170                    &mut self,
171                    threads: usize,
172                    module: &M,
173                    a: &FheUintPrepared<A, $ty, BE>,
174                    b: &FheUintPrepared<B, $ty, BE>,
175                    key: &H,
176                    scratch: &mut Scratch<BE>,
177                ) where
178                    M: ExecuteBDDCircuit2WTo1W<BE>,
179                    A: DataRef,
180                    B: DataRef,
181                    K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
182                    H: GLWEAutomorphismKeyHelper<K, BE>,
183                    Scratch<BE>: ScratchTakeCore<BE>,
184                {
185                    module.execute_bdd_circuit_2w_to_1w_multi_thread(threads, self, &$output_circuits, a, b, key, scratch)
186                }
187            }
188        }
189    };
190}
191
192// a + b
193define_bdd_2w_to_1w_trait!(pub Add, add);
194
195// a - b
196define_bdd_2w_to_1w_trait!(pub Sub, sub);
197
198// a << b
199define_bdd_2w_to_1w_trait!(pub Sll, sll);
200
201// a >> b arithmetic
202define_bdd_2w_to_1w_trait!(pub Sra, sra);
203
204// a >> b logical
205define_bdd_2w_to_1w_trait!(pub Srl, srl);
206
207// signed a < signed b
208define_bdd_2w_to_1w_trait!(pub Slt, slt);
209
210// unsigned a < unsigned b
211define_bdd_2w_to_1w_trait!(pub Sltu, sltu);
212
213// a or b
214define_bdd_2w_to_1w_trait!(pub Or, or);
215
216// a and b
217define_bdd_2w_to_1w_trait!(pub And, and);
218
219// a xor b
220define_bdd_2w_to_1w_trait!(pub Xor, xor);
221
222impl_bdd_2w_to_1w_trait!(
223    Add,
224    add,
225    u32,
226    circuits::u32::add_codegen::AnyBitCircuit,
227    circuits::u32::add_codegen::OUTPUT_CIRCUITS
228);
229
230impl_bdd_2w_to_1w_trait!(
231    Sub,
232    sub,
233    u32,
234    circuits::u32::sub_codegen::AnyBitCircuit,
235    circuits::u32::sub_codegen::OUTPUT_CIRCUITS
236);
237
238impl_bdd_2w_to_1w_trait!(
239    Sll,
240    sll,
241    u32,
242    circuits::u32::sll_codegen::AnyBitCircuit,
243    circuits::u32::sll_codegen::OUTPUT_CIRCUITS
244);
245
246impl_bdd_2w_to_1w_trait!(
247    Sra,
248    sra,
249    u32,
250    circuits::u32::sra_codegen::AnyBitCircuit,
251    circuits::u32::sra_codegen::OUTPUT_CIRCUITS
252);
253
254impl_bdd_2w_to_1w_trait!(
255    Srl,
256    srl,
257    u32,
258    circuits::u32::srl_codegen::AnyBitCircuit,
259    circuits::u32::srl_codegen::OUTPUT_CIRCUITS
260);
261
262impl_bdd_2w_to_1w_trait!(
263    Slt,
264    slt,
265    u32,
266    circuits::u32::slt_codegen::AnyBitCircuit,
267    circuits::u32::slt_codegen::OUTPUT_CIRCUITS
268);
269
270impl_bdd_2w_to_1w_trait!(
271    Sltu,
272    sltu,
273    u32,
274    circuits::u32::sltu_codegen::AnyBitCircuit,
275    circuits::u32::sltu_codegen::OUTPUT_CIRCUITS
276);
277
278impl_bdd_2w_to_1w_trait!(
279    And,
280    and,
281    u32,
282    circuits::u32::and_codegen::AnyBitCircuit,
283    circuits::u32::and_codegen::OUTPUT_CIRCUITS
284);
285
286impl_bdd_2w_to_1w_trait!(
287    Or,
288    or,
289    u32,
290    circuits::u32::or_codegen::AnyBitCircuit,
291    circuits::u32::or_codegen::OUTPUT_CIRCUITS
292);
293
294impl_bdd_2w_to_1w_trait!(
295    Xor,
296    xor,
297    u32,
298    circuits::u32::xor_codegen::AnyBitCircuit,
299    circuits::u32::xor_codegen::OUTPUT_CIRCUITS
300);