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, $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    circuits::u32::add_codegen::AnyBitCircuit,
164    circuits::u32::add_codegen::OUTPUT_CIRCUITS
165);
166
167impl_bdd_2w_to_1w_trait!(
168    Sub,
169    sub,
170    u32,
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    circuits::u32::sll_codegen::AnyBitCircuit,
180    circuits::u32::sll_codegen::OUTPUT_CIRCUITS
181);
182
183impl_bdd_2w_to_1w_trait!(
184    Sra,
185    sra,
186    u32,
187    circuits::u32::sra_codegen::AnyBitCircuit,
188    circuits::u32::sra_codegen::OUTPUT_CIRCUITS
189);
190
191impl_bdd_2w_to_1w_trait!(
192    Srl,
193    srl,
194    u32,
195    circuits::u32::srl_codegen::AnyBitCircuit,
196    circuits::u32::srl_codegen::OUTPUT_CIRCUITS
197);
198
199impl_bdd_2w_to_1w_trait!(
200    Slt,
201    slt,
202    u32,
203    circuits::u32::slt_codegen::AnyBitCircuit,
204    circuits::u32::slt_codegen::OUTPUT_CIRCUITS
205);
206
207impl_bdd_2w_to_1w_trait!(
208    Sltu,
209    sltu,
210    u32,
211    circuits::u32::sltu_codegen::AnyBitCircuit,
212    circuits::u32::sltu_codegen::OUTPUT_CIRCUITS
213);
214
215impl_bdd_2w_to_1w_trait!(
216    And,
217    and,
218    u32,
219    circuits::u32::and_codegen::AnyBitCircuit,
220    circuits::u32::and_codegen::OUTPUT_CIRCUITS
221);
222
223impl_bdd_2w_to_1w_trait!(
224    Or,
225    or,
226    u32,
227    circuits::u32::or_codegen::AnyBitCircuit,
228    circuits::u32::or_codegen::OUTPUT_CIRCUITS
229);
230
231impl_bdd_2w_to_1w_trait!(
232    Xor,
233    xor,
234    u32,
235    circuits::u32::xor_codegen::AnyBitCircuit,
236    circuits::u32::xor_codegen::OUTPUT_CIRCUITS
237);