feat: add piecewise function and function validation

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`

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,22 +534,24 @@ impl<'s> ModelInfo<'s> {
                 self.check_expr(&monop.child);
             }
             AstKind::Call(call) => {
-                // check name in allowed functions
-                let functions = [
-                    "sin", "cos", "tan", "pow", "exp", "log", "sqrt", "abs", "interp1d",
+                // 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) {
-                    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() {

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<Self, ValidationErrors> {
         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!(

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],

diffsl/src/execution/cranelift/codegen.rs

@@ -1248,6 +1248,25 @@ impl CodegenModuleCompile for CraneliftModule<JITModule> {
                 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::<JITModule>::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,45 @@ 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<Value> {
+        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<Value> {
         match &expr.kind {
             AstKind::Integer(value) => Ok(self.builder.ins().iconst(self.int_type, *value)),
@@ -1828,6 +1889,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));

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,