1use std::cmp::{max, min};
  2
  3use crate::{Backend, IntegerType, Module, ZnxBasics, ZnxLayout, ffi::module::MODULE};
  4
  5pub(crate) fn znx_post_process_ternary_op<C, A, B, const NEGATE: bool>(c: &mut C, a: &A, b: &B)
  6where
  7    C: ZnxBasics + ZnxLayout,
  8    A: ZnxBasics + ZnxLayout<Scalar = C::Scalar>,
  9    B: ZnxBasics + ZnxLayout<Scalar = C::Scalar>,
 10    C::Scalar: IntegerType,
 11{
 12    #[cfg(debug_assertions)]
 13    {
 14        assert_ne!(a.as_ptr(), b.as_ptr());
 15        assert_ne!(b.as_ptr(), c.as_ptr());
 16        assert_ne!(a.as_ptr(), c.as_ptr());
 17    }
 18
 19    let a_cols: usize = a.cols();
 20    let b_cols: usize = b.cols();
 21    let c_cols: usize = c.cols();
 22
 23    let min_ab_cols: usize = min(a_cols, b_cols);
 24    let max_ab_cols: usize = max(a_cols, b_cols);
 25
 26    // Copies shared shared cols between (c, max(a, b))
 27    if a_cols != b_cols {
 28        if a_cols > b_cols {
 29            let min_size = min(c.size(), a.size());
 30            (min_ab_cols..min(max_ab_cols, c_cols)).for_each(|i| {
 31                (0..min_size).for_each(|j| {
 32                    c.at_poly_mut(i, j).copy_from_slice(a.at_poly(i, j));
 33                    if NEGATE {
 34                        c.at_poly_mut(i, j).iter_mut().for_each(|x| *x = -*x);
 35                    }
 36                });
 37                (min_size..c.size()).for_each(|j| {
 38                    c.zero_at(i, j);
 39                });
 40            });
 41        } else {
 42            let min_size = min(c.size(), b.size());
 43            (min_ab_cols..min(max_ab_cols, c_cols)).for_each(|i| {
 44                (0..min_size).for_each(|j| {
 45                    c.at_poly_mut(i, j).copy_from_slice(b.at_poly(i, j));
 46                    if NEGATE {
 47                        c.at_poly_mut(i, j).iter_mut().for_each(|x| *x = -*x);
 48                    }
 49                });
 50                (min_size..c.size()).for_each(|j| {
 51                    c.zero_at(i, j);
 52                });
 53            });
 54        }
 55    }
 56
 57    // Zeroes the cols of c > max(a, b).
 58    if c_cols > max_ab_cols {
 59        (max_ab_cols..c_cols).for_each(|i| {
 60            (0..c.size()).for_each(|j| {
 61                c.zero_at(i, j);
 62            })
 63        });
 64    }
 65}
 66
 67#[inline(always)]
 68pub fn apply_binary_op<BE, C, A, B, const NEGATE: bool>(
 69    module: &Module<BE>,
 70    c: &mut C,
 71    a: &A,
 72    b: &B,
 73    op: impl Fn(&mut [C::Scalar], &[A::Scalar], &[B::Scalar]),
 74) where
 75    BE: Backend,
 76    C: ZnxBasics + ZnxLayout,
 77    A: ZnxBasics + ZnxLayout<Scalar = C::Scalar>,
 78    B: ZnxBasics + ZnxLayout<Scalar = C::Scalar>,
 79    C::Scalar: IntegerType,
 80{
 81    #[cfg(debug_assertions)]
 82    {
 83        assert_eq!(a.n(), module.n());
 84        assert_eq!(b.n(), module.n());
 85        assert_eq!(c.n(), module.n());
 86        assert_ne!(a.as_ptr(), b.as_ptr());
 87    }
 88    let a_cols: usize = a.cols();
 89    let b_cols: usize = b.cols();
 90    let c_cols: usize = c.cols();
 91    let min_ab_cols: usize = min(a_cols, b_cols);
 92    let min_cols: usize = min(c_cols, min_ab_cols);
 93    // Applies over shared cols between (a, b, c)
 94    (0..min_cols).for_each(|i| op(c.at_poly_mut(i, 0), a.at_poly(i, 0), b.at_poly(i, 0)));
 95    // Copies/Negates/Zeroes the remaining cols if op is not inplace.
 96    if c.as_ptr() != a.as_ptr() && c.as_ptr() != b.as_ptr() {
 97        znx_post_process_ternary_op::<C, A, B, NEGATE>(c, a, b);
 98    }
 99}
100
101#[inline(always)]
102pub fn apply_unary_op<B: Backend, T: ZnxBasics + ZnxLayout>(
103    module: &Module<B>,
104    b: &mut T,
105    a: &T,
106    op: impl Fn(&mut [T::Scalar], &[T::Scalar]),
107) where
108    <T as ZnxLayout>::Scalar: IntegerType,
109{
110    #[cfg(debug_assertions)]
111    {
112        assert_eq!(a.n(), module.n());
113        assert_eq!(b.n(), module.n());
114    }
115    let a_cols: usize = a.cols();
116    let b_cols: usize = b.cols();
117    let min_cols: usize = min(a_cols, b_cols);
118    // Applies over the shared cols between (a, b)
119    (0..min_cols).for_each(|i| op(b.at_poly_mut(i, 0), a.at_poly(i, 0)));
120    // Zeroes the remaining cols of b.
121    (min_cols..b_cols).for_each(|i| (0..b.size()).for_each(|j| b.zero_at(i, j)));
122}
123
124pub fn ffi_ternary_op_factory<T>(
125    module_ptr: *const MODULE,
126    c_size: usize,
127    c_sl: usize,
128    a_size: usize,
129    a_sl: usize,
130    b_size: usize,
131    b_sl: usize,
132    op_fn: unsafe extern "C" fn(*const MODULE, *mut T, u64, u64, *const T, u64, u64, *const T, u64, u64),
133) -> impl Fn(&mut [T], &[T], &[T]) {
134    move |cv: &mut [T], av: &[T], bv: &[T]| unsafe {
135        op_fn(
136            module_ptr,
137            cv.as_mut_ptr(),
138            c_size as u64,
139            c_sl as u64,
140            av.as_ptr(),
141            a_size as u64,
142            a_sl as u64,
143            bv.as_ptr(),
144            b_size as u64,
145            b_sl as u64,
146        )
147    }
148}
149
150pub fn ffi_binary_op_factory_type_0<T>(
151    module_ptr: *const MODULE,
152    b_size: usize,
153    b_sl: usize,
154    a_size: usize,
155    a_sl: usize,
156    op_fn: unsafe extern "C" fn(*const MODULE, *mut T, u64, u64, *const T, u64, u64),
157) -> impl Fn(&mut [T], &[T]) {
158    move |bv: &mut [T], av: &[T]| unsafe {
159        op_fn(
160            module_ptr,
161            bv.as_mut_ptr(),
162            b_size as u64,
163            b_sl as u64,
164            av.as_ptr(),
165            a_size as u64,
166            a_sl as u64,
167        )
168    }
169}
170
171pub fn ffi_binary_op_factory_type_1<T>(
172    module_ptr: *const MODULE,
173    k: i64,
174    b_size: usize,
175    b_sl: usize,
176    a_size: usize,
177    a_sl: usize,
178    op_fn: unsafe extern "C" fn(*const MODULE, i64, *mut T, u64, u64, *const T, u64, u64),
179) -> impl Fn(&mut [T], &[T]) {
180    move |bv: &mut [T], av: &[T]| unsafe {
181        op_fn(
182            module_ptr,
183            k,
184            bv.as_mut_ptr(),
185            b_size as u64,
186            b_sl as u64,
187            av.as_ptr(),
188            a_size as u64,
189            a_sl as u64,
190        )
191    }
192}