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