mech-math 0.3.3

#[macro_use]
use crate::*;
use super::*;
use num_traits::*;
#[cfg(feature = "matrix")]
use mech_core::matrix::Matrix;

// Add Assign -----------------------------------------------------------------

#[macro_export]
macro_rules! impl_add_assign_match_arms {
  ($fxn_name:ident,$macro_name:ident, $arg:expr) => {
    paste!{
      [<impl_set_ $macro_name _match_arms>]!(
        $fxn_name,
        $arg,
        U8, "u8";
        U16, "u16";
        U32, "u32";
        U64, "u64";
        U128, "u128";
        I8, "i8";
        I16, "i16";
        I32, "i32";
        I64, "i64";
        U128, "u128";
        F32, "f32"; 
        F64, "f64" ;
        C64, "complex";
        R64, "rational";
      )
    }
  }
}

#[cfg(feature = "matrix")]
macro_rules! impl_add_assign_range_fxn_s {
  ($struct_name:ident, $op:ident, $ix:ty) => {
    impl_op_assign_range_fxn_s!($struct_name, $op, $ix);
  }
}

#[cfg(feature = "matrix")]
macro_rules! impl_add_assign_range_fxn_v {
  ($struct_name:ident, $op:ident, $ix:ty) => {
    impl_op_assign_range_fxn_v!($struct_name, $op, $ix);
  }
}

// x += 1 ----------------------------------------------------------------------

impl_assign_scalar_scalar!(Add, +=);
impl_assign_vector_vector!(Add, +=);
impl_assign_vector_scalar!(Add, +=);

pub fn add_assign_math_fxn(sink: Value, source: Value) -> MResult<Box<dyn MechFunction>> {
  impl_op_assign_value_match_arms!(
    Add,
    (sink, source),
    U8,  "u8";
    U16, "u16";
    U32, "u32";
    U64, "u64";
    U128, "u128";
    I8,  "i8";
    I16, "i16";
    I32, "i32";
    I64, "i64";
    U128, "u128";
    F32, "f32";
    F64, "f64";
    R64, "rational";
    C64, "complex";
  )
}

pub struct AddAssignMath {}
impl NativeFunctionCompiler for AddAssignMath {
  fn compile(&self, arguments: &Vec<Value>) -> MResult<Box<dyn MechFunction>> {
    if arguments.len() <= 1 {
      return Err(MechError::new(IncorrectNumberOfArguments { expected: 1, found: arguments.len() },None).with_compiler_loc());
    }
    let sink = arguments[0].clone();
    let source = arguments[1].clone();
    match add_assign_math_fxn(sink.clone(),source.clone()) {
      Ok(fxn) => Ok(fxn),
      Err(x) => {
        match (sink,source) {
          (Value::MutableReference(sink),Value::MutableReference(source)) => { add_assign_math_fxn(sink.borrow().clone(),source.borrow().clone()) },
          (sink,Value::MutableReference(source)) => { add_assign_math_fxn(sink.clone(),source.borrow().clone()) },
          (Value::MutableReference(sink),source) => { add_assign_math_fxn(sink.borrow().clone(),source.clone()) },
          (arg1,arg2) => Err(MechError::new(
              UnhandledFunctionArgumentKind2 { arg: (arg1.kind(),arg2.kind()), fxn_name: "math/add-assign".to_string() },
              None
            ).with_compiler_loc()
          ),
        }
      }
    }
  }
}

register_descriptor! {
  FunctionCompilerDescriptor {
    name: "math/add-assign",
    ptr: &AddAssignMath{},
  }
}

// x[1..3] += 1 ----------------------------------------------------------------

macro_rules! add_assign_1d_range {
  ($source:expr, $ix:expr, $sink:expr) => {
    unsafe { 
      for i in 0..($ix).len() {
        ($sink)[($ix)[i] - 1] += ($source).clone();
      }
    }
  };}

macro_rules! add_assign_1d_range_b {
  ($source:expr, $ix:expr, $sink:expr) => {
    unsafe { 
      for i in 0..($ix).len() {
        if $ix[i] == true {
          ($sink)[i] += ($source).clone();
        }
      }
    }
  };}  

macro_rules! add_assign_1d_range_vec {
  ($source:expr, $ix:expr, $sink:expr) => {
    unsafe { 
      for i in 0..($ix).len() {
        ($sink)[($ix)[i] - 1] += ($source)[i].clone();
      }
    }
  };}

macro_rules! add_assign_1d_range_vec_b {
  ($source:expr, $ix:expr, $sink:expr) => {
    unsafe { 
      for i in 0..($ix).len() {
        if $ix[i] == true {
          ($sink)[i] += ($source)[i].clone();
        }
      }
    }
  };}

#[cfg(feature = "matrix")]
impl_add_assign_range_fxn_s!(AddAssign1DRS,add_assign_1d_range,usize);
#[cfg(feature = "matrix")]
impl_add_assign_range_fxn_s!(AddAssign1DRB,add_assign_1d_range_b,bool);
#[cfg(feature = "matrix")]
impl_add_assign_range_fxn_v!(AddAssign1DRV,add_assign_1d_range_vec,usize);
#[cfg(feature = "matrix")]
impl_add_assign_range_fxn_v!(AddAssign1DRVB,add_assign_1d_range_vec_b,bool);

op_assign_range_fxn!(add_assign_range_fxn, AddAssign1DR);

pub struct AddAssignRange {}
impl NativeFunctionCompiler for AddAssignRange {
  fn compile(&self, arguments: &Vec<Value>) -> MResult<Box<dyn MechFunction>> {
    if arguments.len() <= 1 {
      return Err(MechError::new(IncorrectNumberOfArguments { expected: 1, found: arguments.len() },None).with_compiler_loc());
    }
    let sink: Value = arguments[0].clone();
    let source: Value = arguments[1].clone();
    let ixes = arguments.clone().split_off(2);
    match add_assign_range_fxn(sink.clone(),source.clone(),ixes.clone()) {
      Ok(fxn) => Ok(fxn),
      Err(x) => {
        match (&sink,&ixes,&source) {
          (Value::MutableReference(sink),ixes,Value::MutableReference(source)) => { add_assign_range_fxn(sink.borrow().clone(),source.borrow().clone(),ixes.clone()) },
          (sink,ixes,Value::MutableReference(source)) => { add_assign_range_fxn(sink.clone(),source.borrow().clone(),ixes.clone()) },
          (Value::MutableReference(sink),ixes,source) => { add_assign_range_fxn(sink.borrow().clone(),source.clone(),ixes.clone()) },
          (sink, ixes, source) => Err(MechError::new(
              UnhandledFunctionArgumentIxes { arg: (sink.kind(), ixes.iter().map(|x| x.kind()).collect(), source.kind()), fxn_name: "math/add-assign/range".to_string() },
              None
            ).with_compiler_loc()
          ),
        }
      }
    }
  }
}

register_descriptor! {
  FunctionCompilerDescriptor {
    name: "math/add-assign/range",
    ptr: &AddAssignRange{},
  }
}

// x[1..3,:] += 1 ------------------------------------------------------------------

macro_rules! add_assign_2d_vector_all {
  ($source:expr, $ix:expr, $sink:expr) => {
      for cix in 0..($sink).ncols() {
        for rix in $ix.iter() {
          ($sink).column_mut(cix)[rix - 1] += ($source).clone();
        }
      }
    };}

macro_rules! add_assign_2d_vector_all_b {
  ($source:expr, $ix:expr, $sink:expr) => {
    for cix in 0..($sink).ncols() {
      for rix in 0..$ix.len() {
        if $ix[rix] == true {
          ($sink).column_mut(cix)[rix] += ($source).clone();
        }
      }
    }
  };} 

macro_rules! add_assign_2d_vector_all_mat {
  ($source:expr, $ix:expr, $sink:expr) => {
    {
      let nsrc = $source.nrows();
      for (i, &rix) in $ix.iter().enumerate() {
        let row_index = rix - 1;
        let mut sink_row = $sink.row_mut(row_index);
        let src_row = $source.row(i % nsrc); // wrap around!
        for (dst, src) in sink_row.iter_mut().zip(src_row.iter()) {
          *dst += *src;
        }
      }
    }
  };}

macro_rules! add_assign_2d_vector_all_mat_b {
  ($source:expr, $ix:expr, $sink:expr) => {
    {
      let mut src_i = 0;
      for (i, rix) in (&$ix).iter().enumerate() {
        if *rix == true {
          let mut sink_row = ($sink).row_mut(i);
          let src_row = ($source).row(src_i);
          for (dst, src) in sink_row.iter_mut().zip(src_row.iter()) {
            *dst += *src;
          }
          src_i += 1;
        }
      }
    }
  };}

#[cfg(feature = "matrix")]
impl_add_assign_range_fxn_s!(AddAssign2DRAS, add_assign_2d_vector_all,usize);
#[cfg(feature = "matrix")]
impl_add_assign_range_fxn_s!(AddAssign2DRASB,add_assign_2d_vector_all_b,bool);
#[cfg(feature = "matrix")]
impl_add_assign_range_fxn_v!(AddAssign2DRAV, add_assign_2d_vector_all_mat,usize);
#[cfg(feature = "matrix")]
impl_add_assign_range_fxn_v!(AddAssign2DRAVB,add_assign_2d_vector_all_mat_b,bool);

op_assign_range_all_fxn!(add_assign_range_all_fxn, AddAssign2DRA);

pub struct AddAssignRangeAll {}
impl NativeFunctionCompiler for AddAssignRangeAll {
  fn compile(&self, arguments: &Vec<Value>) -> MResult<Box<dyn MechFunction>> {
    if arguments.len() <= 1 {
      return Err(MechError::new(IncorrectNumberOfArguments { expected: 1, found: arguments.len() },None).with_compiler_loc());
    }
    let sink: Value = arguments[0].clone();
    let source: Value = arguments[1].clone();
    let ixes = arguments.clone().split_off(2);
    match add_assign_range_all_fxn(sink.clone(),source.clone(),ixes.clone()) {
      Ok(fxn) => Ok(fxn),
      Err(_) => {
        match (sink,ixes,source) {
          (Value::MutableReference(sink),ixes,Value::MutableReference(source)) => { add_assign_range_all_fxn(sink.borrow().clone(),source.borrow().clone(),ixes.clone()) },
          (sink,ixes,Value::MutableReference(source)) => { add_assign_range_all_fxn(sink.clone(),source.borrow().clone(),ixes.clone()) },
          (Value::MutableReference(sink),ixes,source) => { add_assign_range_all_fxn(sink.borrow().clone(),source.clone(),ixes.clone()) },
          (sink, ixes, source) => Err(MechError::new(
              UnhandledFunctionArgumentIxes { arg: (sink.kind(), ixes.iter().map(|x| x.kind()).collect(), source.kind()), fxn_name: "math/add-assign/range-all".to_string() },
              None
            ).with_compiler_loc()
          ),
        }
      }
    }
  }
}

register_descriptor! {
  FunctionCompilerDescriptor {
    name: "math/add-assign/range-all",
    ptr: &AddAssignRangeAll{},
  }
}