1use std::cmp::{max, min};
 2
 3use crate::{Backend, IntegerType, Module, ZnxBasics, ZnxLayout, ffi::module::MODULE};
 4
 5#[inline(always)]
 6pub fn apply_unary_op<B: Backend, T: ZnxBasics + ZnxLayout>(
 7    module: &Module<B>,
 8    b: &mut T,
 9    a: &T,
10    op: impl Fn(&mut [T::Scalar], &[T::Scalar]),
11) where
12    <T as ZnxLayout>::Scalar: IntegerType,
13{
14    #[cfg(debug_assertions)]
15    {
16        assert_eq!(a.n(), module.n());
17        assert_eq!(b.n(), module.n());
18    }
19    let a_cols: usize = a.cols();
20    let b_cols: usize = b.cols();
21    let min_cols: usize = min(a_cols, b_cols);
22    // Applies over the shared cols between (a, b)
23    (0..min_cols).for_each(|i| op(b.at_poly_mut(i, 0), a.at_poly(i, 0)));
24    // Zeroes the remaining cols of b.
25    (min_cols..b_cols).for_each(|i| (0..b.size()).for_each(|j| b.zero_at(i, j)));
26}
27
28pub fn ffi_ternary_op_factory<T>(
29    module_ptr: *const MODULE,
30    c_size: usize,
31    c_sl: usize,
32    a_size: usize,
33    a_sl: usize,
34    b_size: usize,
35    b_sl: usize,
36    op_fn: unsafe extern "C" fn(*const MODULE, *mut T, u64, u64, *const T, u64, u64, *const T, u64, u64),
37) -> impl Fn(&mut [T], &[T], &[T]) {
38    move |cv: &mut [T], av: &[T], bv: &[T]| unsafe {
39        op_fn(
40            module_ptr,
41            cv.as_mut_ptr(),
42            c_size as u64,
43            c_sl as u64,
44            av.as_ptr(),
45            a_size as u64,
46            a_sl as u64,
47            bv.as_ptr(),
48            b_size as u64,
49            b_sl as u64,
50        )
51    }
52}
53
54pub fn ffi_binary_op_factory_type_0<T>(
55    module_ptr: *const MODULE,
56    b_size: usize,
57    b_sl: usize,
58    a_size: usize,
59    a_sl: usize,
60    op_fn: unsafe extern "C" fn(*const MODULE, *mut T, u64, u64, *const T, u64, u64),
61) -> impl Fn(&mut [T], &[T]) {
62    move |bv: &mut [T], av: &[T]| unsafe {
63        op_fn(
64            module_ptr,
65            bv.as_mut_ptr(),
66            b_size as u64,
67            b_sl as u64,
68            av.as_ptr(),
69            a_size as u64,
70            a_sl as u64,
71        )
72    }
73}
74
75pub fn ffi_binary_op_factory_type_1<T>(
76    module_ptr: *const MODULE,
77    k: i64,
78    b_size: usize,
79    b_sl: usize,
80    a_size: usize,
81    a_sl: usize,
82    op_fn: unsafe extern "C" fn(*const MODULE, i64, *mut T, u64, u64, *const T, u64, u64),
83) -> impl Fn(&mut [T], &[T]) {
84    move |bv: &mut [T], av: &[T]| unsafe {
85        op_fn(
86            module_ptr,
87            k,
88            bv.as_mut_ptr(),
89            b_size as u64,
90            b_sl as u64,
91            av.as_ptr(),
92            a_size as u64,
93            a_sl as u64,
94        )
95    }
96}