From e50abe69c339480dc21dd2a2df6c9c2f419bedfd Mon Sep 17 00:00:00 2001 From: martinjrobins Date: Sat, 18 Jul 2026 07:34:26 +0000 Subject: [PATCH 1/3] feat: add piecewise function --- diffsl/src/continuous/builder.rs | 33 +++++--- diffsl/src/execution/compiler.rs | 17 ++++ diffsl/src/execution/cranelift/codegen.rs | 65 +++++++++++++++ diffsl/src/execution/data_layout.rs | 9 +++ diffsl/src/execution/functions.rs | 99 +++++++++++++++++++++++ diffsl/src/execution/llvm/codegen.rs | 35 ++++++++ 6 files changed, 248 insertions(+), 10 deletions(-) diff --git a/diffsl/src/continuous/builder.rs b/diffsl/src/continuous/builder.rs index f2a2588..e2feb77 100644 --- a/diffsl/src/continuous/builder.rs +++ b/diffsl/src/continuous/builder.rs @@ -536,18 +536,31 @@ impl<'s> ModelInfo<'s> { // check name in allowed functions let functions = [ "sin", "cos", "tan", "pow", "exp", "log", "sqrt", "abs", "interp1d", + "piecewise", ]; if functions.contains(&call.fn_name) { - let expected_nargs = match call.fn_name { - "interp1d" => 3, - "pow" => 2, - _ => 1, - }; - if call.args.len() != expected_nargs { - self.errors.push(Output::new( - format!("incorrect number of given arguments ({} instead of {}) for function {}", call.args.len(), expected_nargs, call.fn_name), - expr.span, - )); + if call.fn_name == "piecewise" { + if call.args.len() < 3 || call.args.len() % 2 == 0 { + self.errors.push(Output::new( + format!( + "piecewise requires an odd number of arguments >= 3, got {}", + call.args.len() + ), + expr.span, + )); + } + } else { + let expected_nargs = match call.fn_name { + "interp1d" => 3, + "pow" => 2, + _ => 1, + }; + if call.args.len() != expected_nargs { + self.errors.push(Output::new( + format!("incorrect number of given arguments ({} instead of {}) for function {}", call.args.len(), expected_nargs, call.fn_name), + expr.span, + )); + } } for arg in call.args.iter() { if let AstKind::CallArg(call_arg) = &arg.kind { diff --git a/diffsl/src/execution/compiler.rs b/diffsl/src/execution/compiler.rs index 58f5f46..61dcec3 100644 --- a/diffsl/src/execution/compiler.rs +++ b/diffsl/src/execution/compiler.rs @@ -2940,6 +2940,23 @@ mod tests { " expect "r" vec![1.0], } + tensor_test! { + piecewise_single: "r { piecewise(5.0, 42.0, 0.0) }" expect "r" vec![42.0], + piecewise_first_true: "r { piecewise(1.0, 10.0, -1.0, 20.0, 30.0) }" expect "r" vec![10.0], + piecewise_second_true: "r { piecewise(-1.0, 10.0, 1.0, 20.0, 30.0) }" expect "r" vec![20.0], + piecewise_fallback: "r { piecewise(-1.0, 10.0, -2.0, 20.0, 30.0) }" expect "r" vec![30.0], + piecewise_zero_check: "r { piecewise(0.0, 100.0, 200.0) }" expect "r" vec![100.0], + piecewise_expr: "r { piecewise(-3.0, 1.0 + 2.0, 3.0, 3.0 * 4.0, 5.0 + 6.0) }" expect "r" vec![12.0], + piecewise_heaviside_equiv: " + x { 0.5 } + r { piecewise(x - 1.0, x * x, -x) } + " expect "r" vec![-0.5], + piecewise_var_check: " + x { 0.5 } + r { piecewise(x, 100.0, 200.0) } + " expect "r" vec![100.0], + } + tensor_test! { scalar_vec_index_bug: "a_i { b = 1 } c_i { (0:2): 0, (2:3): 0.5 } d_i { (0:1): 2 } r_i { c_i * (d * b), c_i * (d_i * b) }" expect "r" vec![0.0, 0.0, 1.0, 0.0, 0.0, 1.0], sparse_mat_mul12: "A_ij { (0,0):1, (0,1):1, (1,1):1, (2,2):1, (3,3):1, (4,4):1, (5,5):1 } b2_i { (1:6): 1 } r_i { A_ij * b2_j }" expect "r" vec![1.0, 1.0, 1.0, 1.0, 1.0, 1.0], diff --git a/diffsl/src/execution/cranelift/codegen.rs b/diffsl/src/execution/cranelift/codegen.rs index dcb5b36..0ab625b 100644 --- a/diffsl/src/execution/cranelift/codegen.rs +++ b/diffsl/src/execution/cranelift/codegen.rs @@ -1248,6 +1248,25 @@ impl CodegenModuleCompile for CraneliftModule { df_ptr, ); } + for i in 0..crate::execution::functions::PIECEWISE_FUNCTIONS.len() { + let (f_name, f_ptr, df_ptr) = match real_type { + RealType::F32 => ( + crate::execution::functions::PIECEWISE_FUNCTIONS_F32[i].0, + crate::execution::functions::PIECEWISE_FUNCTIONS_F32[i].1 as *const u8, + crate::execution::functions::PIECEWISE_FUNCTIONS_F32[i].2 as *const u8, + ), + RealType::F64 => ( + crate::execution::functions::PIECEWISE_FUNCTIONS_F64[i].0, + crate::execution::functions::PIECEWISE_FUNCTIONS_F64[i].1 as *const u8, + crate::execution::functions::PIECEWISE_FUNCTIONS_F64[i].2 as *const u8, + ), + }; + builder.symbol(f_name, f_ptr); + builder.symbol( + CraneliftCodeGen::::get_function_name(f_name, true), + df_ptr, + ); + } let module = JITModule::new(builder); Self::new( @@ -1389,6 +1408,9 @@ impl<'ctx, M: Module> CraneliftCodeGen<'ctx, M> { AstKind::Call(call) if call.fn_name == "interp1d" => { self.jit_compile_interp1d(name, call, index, elmt, expr_index, expr.span) } + AstKind::Call(call) if call.fn_name == "piecewise" => { + self.jit_compile_piecewise(name, call, index, elmt, expr_index) + } AstKind::Call(call) => match self.get_function(call.fn_name, call.is_tangent) { Some(function) => { let mut args = Vec::new(); @@ -1753,6 +1775,46 @@ impl<'ctx, M: Module> CraneliftCodeGen<'ctx, M> { Ok(ret_value) } + fn jit_compile_piecewise( + &mut self, + name: &str, + call: &crate::ast::Call, + index: &[Value], + elmt: &TensorBlock, + expr_index: Value, + ) -> Result { + let is_tangent = call.is_tangent; + let n = call.args.len(); + + let mut args = Vec::new(); + for arg in call.args.iter() { + let arg_val = + self.jit_compile_expr(name, arg.as_ref(), index, elmt, expr_index)?; + args.push(arg_val); + } + + let arg_size = (n as u32) * self.real_type.bytes(); + let ss_data = StackSlotData::new(StackSlotKind::ExplicitSlot, arg_size, 0); + let ss = self.builder.create_sized_stack_slot(ss_data); + + let byte_size = self.real_type.bytes() as i32; + for (i, arg) in args.iter().enumerate() { + self.builder + .ins() + .stack_store(*arg, ss, (i as i32) * byte_size); + } + + let ptr = self.builder.ins().stack_addr(self.real_ptr_type, ss, 0); + let count = self.builder.ins().iconst(self.int_type, n as i64); + + let function = self + .get_function("piecewise_impl", is_tangent) + .ok_or_else(|| anyhow!("piecewise_impl not found"))?; + + let call_inst = self.builder.ins().call(function, &[ptr, count]); + Ok(self.builder.inst_results(call_inst)[0]) + } + fn jit_compile_integer_expr(&mut self, expr: &Ast) -> Result { match &expr.kind { AstKind::Integer(value) => Ok(self.builder.ins().iconst(self.int_type, *value)), @@ -1828,6 +1890,9 @@ impl<'ctx, M: Module> CraneliftCodeGen<'ctx, M> { if is_tangent { sig.params.push(AbiParam::new(self.real_type)); } + } else if base_name == "piecewise_impl" { + sig.params.push(AbiParam::new(self.real_ptr_type)); + sig.params.push(AbiParam::new(self.int_type)); } else { for _ in 0..num_args { sig.params.push(AbiParam::new(self.real_type)); diff --git a/diffsl/src/execution/data_layout.rs b/diffsl/src/execution/data_layout.rs index 789c99e..0cb95f4 100644 --- a/diffsl/src/execution/data_layout.rs +++ b/diffsl/src/execution/data_layout.rs @@ -575,6 +575,15 @@ impl DataLayout { ("pow", [x, y]) => x.powf(*y), ("min", [x, y]) => x.min(*y), ("max", [x, y]) => x.max(*y), + ("piecewise", args) => { + let num_branches = (args.len() - 1) / 2; + for i in 0..num_branches { + if args[2 * i] >= 0.0 { + return args[2 * i + 1]; + } + } + args[args.len() - 1] + } _ => panic!( "unknown constant function call '{}' with {} args", name, diff --git a/diffsl/src/execution/functions.rs b/diffsl/src/execution/functions.rs index e703244..3c42b6b 100644 --- a/diffsl/src/execution/functions.rs +++ b/diffsl/src/execution/functions.rs @@ -13,6 +13,8 @@ type Interp1dFnF64 = extern "C" fn(*const f64, usize, usize, usize, f64) -> f64; type Interp1dFnF32 = extern "C" fn(*const f32, usize, usize, usize, f32) -> f32; type Interp1dGradFnF64 = extern "C" fn(*const f64, usize, usize, usize, f64, f64) -> f64; type Interp1dGradFnF32 = extern "C" fn(*const f32, usize, usize, usize, f32, f32) -> f32; +type PiecewiseFnF64 = extern "C" fn(*const f64, usize) -> f64; +type PiecewiseFnF32 = extern "C" fn(*const f32, usize) -> f32; pub fn function_symbol_name(base_name: &str, real_type: RealType, is_tangent: bool) -> String { let suffix = real_type.as_str(); @@ -77,6 +79,18 @@ fn resolve_by_real_type(name: &str, is_tangent: bool, real_type: &RealType) -> O *f as *const u8 } }) + }) + .or_else(|| { + PIECEWISE_FUNCTIONS_F64 + .iter() + .find(|(n, _, _)| *n == name) + .map(|(_, f, df)| { + if is_tangent { + *df as *const u8 + } else { + *f as *const u8 + } + }) }), RealType::F32 => FUNCTIONS_F32 .iter() @@ -111,6 +125,18 @@ fn resolve_by_real_type(name: &str, is_tangent: bool, real_type: &RealType) -> O *f as *const u8 } }) + }) + .or_else(|| { + PIECEWISE_FUNCTIONS_F32 + .iter() + .find(|(n, _, _)| *n == name) + .map(|(_, f, df)| { + if is_tangent { + *df as *const u8 + } else { + *f as *const u8 + } + }) }), } } @@ -179,6 +205,15 @@ pub const INTERP1D_FUNCTIONS_F32: &[(&str, Interp1dFnF32, Interp1dGradFnF32)] = pub const INTERP1D_FUNCTIONS: &[(&str, Interp1dFnF64, Interp1dGradFnF64)] = INTERP1D_FUNCTIONS_F64; +pub const PIECEWISE_FUNCTIONS_F64: &[(&str, PiecewiseFnF64, PiecewiseFnF64)] = + &[("piecewise_impl", piecewise_impl_f64, dpiecewise_impl_f64)]; + +pub const PIECEWISE_FUNCTIONS_F32: &[(&str, PiecewiseFnF32, PiecewiseFnF32)] = + &[("piecewise_impl", piecewise_impl_f32, dpiecewise_impl_f32)]; + +pub const PIECEWISE_FUNCTIONS: &[(&str, PiecewiseFnF64, PiecewiseFnF64)] = + PIECEWISE_FUNCTIONS_F64; + pub fn function_resolver(name: &str) -> Option<*const u8> { let (base_name, is_tangent, real_type) = parse_function_name(name); @@ -293,6 +328,16 @@ pub fn function_num_args(name: &str, is_tangent: bool) -> Option { return Some(5 * multiplier); } + if PIECEWISE_FUNCTIONS_F64 + .iter() + .any(|(n, _, _)| n == &base_name) + || PIECEWISE_FUNCTIONS_F32 + .iter() + .any(|(n, _, _)| n == &base_name) + { + return Some(2 * multiplier); + } + None } @@ -721,3 +766,57 @@ extern "C" fn dinterp1d_impl_f32( dq * slope } } + +extern "C" fn piecewise_impl_f64(args: *const f64, n: usize) -> f64 { + unsafe { + let args = std::slice::from_raw_parts(args, n); + let num_branches = (n - 1) / 2; + for i in 0..num_branches { + if args[2 * i] >= 0.0 { + return args[2 * i + 1]; + } + } + args[n - 1] + } +} + +extern "C" fn dpiecewise_impl_f64(args: *const f64, n: usize) -> f64 { + unsafe { + let args = std::slice::from_raw_parts(args, n); + let total_primal = n / 2; + let num_branches = (total_primal - 1) / 2; + for i in 0..num_branches { + if args[4 * i] >= 0.0 { + return args[4 * i + 3]; + } + } + args[n - 1] + } +} + +extern "C" fn piecewise_impl_f32(args: *const f32, n: usize) -> f32 { + unsafe { + let args = std::slice::from_raw_parts(args, n); + let num_branches = (n - 1) / 2; + for i in 0..num_branches { + if args[2 * i] >= 0.0 { + return args[2 * i + 1]; + } + } + args[n - 1] + } +} + +extern "C" fn dpiecewise_impl_f32(args: *const f32, n: usize) -> f32 { + unsafe { + let args = std::slice::from_raw_parts(args, n); + let total_primal = n / 2; + let num_branches = (total_primal - 1) / 2; + for i in 0..num_branches { + if args[4 * i] >= 0.0 { + return args[4 * i + 3]; + } + } + args[n - 1] + } +} diff --git a/diffsl/src/execution/llvm/codegen.rs b/diffsl/src/execution/llvm/codegen.rs index 8e9e06c..1f600d5 100644 --- a/diffsl/src/execution/llvm/codegen.rs +++ b/diffsl/src/execution/llvm/codegen.rs @@ -3028,6 +3028,9 @@ impl<'ctx> CodeGen<'ctx> { AstKind::Call(call) if call.fn_name == "interp1d" => { self.jit_compile_interp1d(name, call, index, elmt, expr_index, expr.span) } + AstKind::Call(call) if call.fn_name == "piecewise" => { + self.jit_compile_piecewise(name, call, index, elmt, expr_index) + } AstKind::Call(call) => match self.get_function(call.fn_name) { Some(function) => { let mut args: Vec = Vec::new(); @@ -3402,6 +3405,38 @@ impl<'ctx> CodeGen<'ctx> { Ok(ret) } + fn jit_compile_piecewise( + &mut self, + name: &str, + call: &crate::ast::Call, + index: &[IntValue<'ctx>], + elmt: &TensorBlock, + expr_index: IntValue<'ctx>, + ) -> Result> { + let n = call.args.len(); + let num_branches = (n - 1) / 2; + + let fallback = + self.jit_compile_expr(name, &call.args[n - 1], index, elmt, expr_index)?; + + let mut result = fallback; + for i in (0..num_branches).rev() { + let check = + self.jit_compile_expr(name, &call.args[2 * i], index, elmt, expr_index)?; + let value = + self.jit_compile_expr(name, &call.args[2 * i + 1], index, elmt, expr_index)?; + let zero = self.real_type.const_float(0.0); + let cond = self + .builder + .build_float_compare(FloatPredicate::OGE, check, zero, name)?; + result = self + .builder + .build_select(cond, value, result, name)? + .into_float_value(); + } + Ok(result) + } + fn jit_compile_integer_expr(&mut self, expr: &Ast, name: &str) -> Result> { match &expr.kind { AstKind::Integer(value) => Ok(self.int_type.const_int(*value as u64, true)), From bbb4fd7ce4371e4951235efb2c9748d230ec52e9 Mon Sep 17 00:00:00 2001 From: martinjrobins Date: Sat, 18 Jul 2026 08:13:11 +0000 Subject: [PATCH 2/3] fix: add function arg count validation --- book/src/functions.md | 16 +++++++++ diffsl/src/continuous/builder.rs | 34 +++++------------- diffsl/src/discretise/discrete_model.rs | 32 +++++++++++++++++ diffsl/src/execution/functions.rs | 48 +++++++++++++++++++++++++ 4 files changed, 104 insertions(+), 26 deletions(-) diff --git a/book/src/functions.md b/book/src/functions.md index 8c005d6..9083403 100644 --- a/book/src/functions.md +++ b/book/src/functions.md @@ -21,6 +21,7 @@ The DiffSL supports the following mathematical functions that can be used in an * `sigmoid(x)` - sigmoid function of x * `heaviside(x)` - Heaviside step function of x * `interp1d(x_i, y_i, q)` - 1D piecewise linear interpolation +* `piecewise(c0, v0, c1, v1, ..., fallback)` - piecewise conditional function You can use these functions as part of an expression in the DSL. For example, to define a variable `a` that is the sine of another variable `b`, you can write: @@ -29,6 +30,21 @@ b { 1.0 } a { sin(b) } ``` +### `piecewise(c0, v0, c1, v1, ..., fallback)` + +The `piecewise` function takes an odd number of arguments (minimum 3) and evaluates +to a piecewise mathematical function. Arguments are grouped in pairs `(check, value)`, +with a final `fallback`. For each pair, if `check >= 0`, the function returns +`value`. If no check passes, it returns `fallback`. For example, to represent +the piecewise function `f(x) = x^2 for x >= 0, -x for x < 0`:` + +```diffsl +# +x { 0.5 } +r { piecewise(x, x * x, -x) } +# x = 0.5 >= 0 → return x * x = 0.25 +``` + ### `interp1d(x_i, y_i, q)` The `interp1d` function performs 1D piecewise linear interpolation. `x_i` and `y_i` diff --git a/diffsl/src/continuous/builder.rs b/diffsl/src/continuous/builder.rs index e2feb77..27632ac 100644 --- a/diffsl/src/continuous/builder.rs +++ b/diffsl/src/continuous/builder.rs @@ -3,6 +3,7 @@ use crate::ast::Ast; use crate::ast::AstKind; use crate::ast::Model; use crate::ast::StringSpan; +use crate::execution::functions::check_function_args; use pest::Span; use std::cell::RefCell; use std::collections::HashMap; @@ -533,35 +534,16 @@ impl<'s> ModelInfo<'s> { self.check_expr(&monop.child); } AstKind::Call(call) => { - // check name in allowed functions - let functions = [ + // validate arg count for builtin functions + if let Err(msg) = check_function_args(call.fn_name, call.args.len()) { + self.errors.push(Output::new(msg, expr.span)); + } + // keyword arg check for builtins + let builtin = [ "sin", "cos", "tan", "pow", "exp", "log", "sqrt", "abs", "interp1d", "piecewise", ]; - if functions.contains(&call.fn_name) { - if call.fn_name == "piecewise" { - if call.args.len() < 3 || call.args.len() % 2 == 0 { - self.errors.push(Output::new( - format!( - "piecewise requires an odd number of arguments >= 3, got {}", - call.args.len() - ), - expr.span, - )); - } - } else { - let expected_nargs = match call.fn_name { - "interp1d" => 3, - "pow" => 2, - _ => 1, - }; - if call.args.len() != expected_nargs { - self.errors.push(Output::new( - format!("incorrect number of given arguments ({} instead of {}) for function {}", call.args.len(), expected_nargs, call.fn_name), - expr.span, - )); - } - } + if builtin.contains(&call.fn_name) { for arg in call.args.iter() { if let AstKind::CallArg(call_arg) = &arg.kind { if call_arg.name.is_some() { diff --git a/diffsl/src/discretise/discrete_model.rs b/diffsl/src/discretise/discrete_model.rs index 8dc6a46..9d26bb1 100644 --- a/diffsl/src/discretise/discrete_model.rs +++ b/diffsl/src/discretise/discrete_model.rs @@ -408,6 +408,29 @@ impl<'s> DiscreteModel<'s> { } } + fn check_function_args_in_expr(expr: &Ast, env: &mut Env) { + match &expr.kind { + AstKind::Call(call) => { + if let Err(msg) = + crate::execution::functions::check_function_args(call.fn_name, call.args.len()) + { + env.errs_mut().push(ValidationError::new(msg, expr.span)); + } + for arg in &call.args { + Self::check_function_args_in_expr(arg, env); + } + } + AstKind::CallArg(arg) => Self::check_function_args_in_expr(&arg.expression, env), + AstKind::Binop(binop) => { + Self::check_function_args_in_expr(&binop.left, env); + Self::check_function_args_in_expr(&binop.right, env); + } + AstKind::Monop(monop) => Self::check_function_args_in_expr(&monop.child, env), + AstKind::Assignment(a) => Self::check_function_args_in_expr(&a.expr, env), + _ => {} + } + } + pub fn build(name: &'s str, model: &'s ast::DsModel) -> Result { let mut env = Env::new(); if model.has_inputs { @@ -754,6 +777,13 @@ impl<'s> DiscreteModel<'s> { Vec::new() }; + // validate function arg counts across all tensors + for tensor in ret.all_tensors() { + for blk in tensor.elmts() { + Self::check_function_args_in_expr(blk.expr(), &mut env); + } + } + if env.errs().is_empty() { Ok(ret) } else { @@ -1607,6 +1637,8 @@ mod tests { error_divide_by_zero: "a_i { (0): 1, (2): 2 } b_i { (2): 1 } c_i { a_i / b_i }" errors ["divide-by-zero",], error_divide_by_zero2: "a_i { (0:3): 1 } b_i { (2): 1 } c_i { a_i / b_i }" errors ["divide-by-zero",], slice_sparse_vec: "A_i { (0): 1, (2): 3 } B_i { A_i[0:1] }" errors ["can only index dense 1D variables",], + piecewise_even_args: "r { piecewise(0.0, 1.0) }" errors ["requires an odd number of arguments >= 3",], + piecewise_one_arg: "r { piecewise(0.0) }" errors ["requires an odd number of arguments >= 3",], ); tensor_tests!( diff --git a/diffsl/src/execution/functions.rs b/diffsl/src/execution/functions.rs index 3c42b6b..1b6fe30 100644 --- a/diffsl/src/execution/functions.rs +++ b/diffsl/src/execution/functions.rs @@ -341,6 +341,54 @@ pub fn function_num_args(name: &str, is_tangent: bool) -> Option { None } +pub fn check_function_args(name: &str, n_args: usize) -> Result<(), String> { + if FUNCTIONS_F64.iter().any(|(n, _, _)| n == &name) + || FUNCTIONS_F32.iter().any(|(n, _, _)| n == &name) + { + if n_args != 1 { + return Err(format!( + "function '{}' expects 1 argument, got {}", + name, n_args + )); + } + return Ok(()); + } + + if TWO_ARG_FUNCTIONS_F64.iter().any(|(n, _, _)| n == &name) + || TWO_ARG_FUNCTIONS_F32.iter().any(|(n, _, _)| n == &name) + { + if n_args != 2 { + return Err(format!( + "function '{}' expects 2 arguments, got {}", + name, n_args + )); + } + return Ok(()); + } + + if name == "interp1d" { + if n_args != 3 { + return Err(format!( + "function 'interp1d' expects 3 arguments, got {}", + n_args + )); + } + return Ok(()); + } + + if name == "piecewise" { + if n_args < 3 || n_args % 2 == 0 { + return Err(format!( + "function 'piecewise' requires an odd number of arguments >= 3, got {}", + n_args + )); + } + return Ok(()); + } + + Ok(()) +} + // Explicit f64 versions extern "C" fn sin_f64(x: f64) -> f64 { x.sin() From f72525fb32c16dfb823319b08af73b80dff2d8e6 Mon Sep 17 00:00:00 2001 From: martinjrobins Date: Sat, 18 Jul 2026 08:31:53 +0000 Subject: [PATCH 3/3] fix: clippy --- diffsl/src/continuous/builder.rs | 10 +++++++++- diffsl/src/execution/cranelift/codegen.rs | 3 +-- diffsl/src/execution/functions.rs | 5 ++--- diffsl/src/execution/llvm/codegen.rs | 6 ++---- 4 files changed, 14 insertions(+), 10 deletions(-) diff --git a/diffsl/src/continuous/builder.rs b/diffsl/src/continuous/builder.rs index 27632ac..70b3e4c 100644 --- a/diffsl/src/continuous/builder.rs +++ b/diffsl/src/continuous/builder.rs @@ -540,7 +540,15 @@ impl<'s> ModelInfo<'s> { } // keyword arg check for builtins let builtin = [ - "sin", "cos", "tan", "pow", "exp", "log", "sqrt", "abs", "interp1d", + "sin", + "cos", + "tan", + "pow", + "exp", + "log", + "sqrt", + "abs", + "interp1d", "piecewise", ]; if builtin.contains(&call.fn_name) { diff --git a/diffsl/src/execution/cranelift/codegen.rs b/diffsl/src/execution/cranelift/codegen.rs index 0ab625b..e594ab3 100644 --- a/diffsl/src/execution/cranelift/codegen.rs +++ b/diffsl/src/execution/cranelift/codegen.rs @@ -1788,8 +1788,7 @@ impl<'ctx, M: Module> CraneliftCodeGen<'ctx, M> { let mut args = Vec::new(); for arg in call.args.iter() { - let arg_val = - self.jit_compile_expr(name, arg.as_ref(), index, elmt, expr_index)?; + let arg_val = self.jit_compile_expr(name, arg.as_ref(), index, elmt, expr_index)?; args.push(arg_val); } diff --git a/diffsl/src/execution/functions.rs b/diffsl/src/execution/functions.rs index 1b6fe30..0d8f1f1 100644 --- a/diffsl/src/execution/functions.rs +++ b/diffsl/src/execution/functions.rs @@ -211,8 +211,7 @@ pub const PIECEWISE_FUNCTIONS_F64: &[(&str, PiecewiseFnF64, PiecewiseFnF64)] = pub const PIECEWISE_FUNCTIONS_F32: &[(&str, PiecewiseFnF32, PiecewiseFnF32)] = &[("piecewise_impl", piecewise_impl_f32, dpiecewise_impl_f32)]; -pub const PIECEWISE_FUNCTIONS: &[(&str, PiecewiseFnF64, PiecewiseFnF64)] = - PIECEWISE_FUNCTIONS_F64; +pub const PIECEWISE_FUNCTIONS: &[(&str, PiecewiseFnF64, PiecewiseFnF64)] = PIECEWISE_FUNCTIONS_F64; pub fn function_resolver(name: &str) -> Option<*const u8> { let (base_name, is_tangent, real_type) = parse_function_name(name); @@ -377,7 +376,7 @@ pub fn check_function_args(name: &str, n_args: usize) -> Result<(), String> { } if name == "piecewise" { - if n_args < 3 || n_args % 2 == 0 { + if n_args < 3 || n_args.is_multiple_of(2) { return Err(format!( "function 'piecewise' requires an odd number of arguments >= 3, got {}", n_args diff --git a/diffsl/src/execution/llvm/codegen.rs b/diffsl/src/execution/llvm/codegen.rs index 1f600d5..5fb1355 100644 --- a/diffsl/src/execution/llvm/codegen.rs +++ b/diffsl/src/execution/llvm/codegen.rs @@ -3416,13 +3416,11 @@ impl<'ctx> CodeGen<'ctx> { let n = call.args.len(); let num_branches = (n - 1) / 2; - let fallback = - self.jit_compile_expr(name, &call.args[n - 1], index, elmt, expr_index)?; + let fallback = self.jit_compile_expr(name, &call.args[n - 1], index, elmt, expr_index)?; let mut result = fallback; for i in (0..num_branches).rev() { - let check = - self.jit_compile_expr(name, &call.args[2 * i], index, elmt, expr_index)?; + let check = self.jit_compile_expr(name, &call.args[2 * i], index, elmt, expr_index)?; let value = self.jit_compile_expr(name, &call.args[2 * i + 1], index, elmt, expr_index)?; let zero = self.real_type.const_float(0.0);