fix: fix and document wrap_mtkparameters

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -398,50 +398,64 @@ function wrap_array_vars(
     end
 end
 
-function wrap_mtkparameters(sys::AbstractSystem, isscalar::Bool)
+const MTKPARAMETERS_ARG = Sym{Vector{Vector}}(:___mtkparameters___)
+
+"""
+    wrap_mtkparameters(sys::AbstractSystem, isscalar::Bool, p_start = 2)
+
+Return function(s) to be passed to the `wrap_code` keyword of `build_function` which
+allow the compiled function to be called as `f(u, p, t)` where `p isa MTKParameters`
+instead of `f(u, p..., t)`. `isscalar` denotes whether the function expression being
+wrapped is for a scalar value. `p_start` is the index of the argument containing
+the first parameter vector in the out-of-place version of the function. For example,
+if a history function (DDEs) was passed before `p`, then the function before wrapping
+would have the signature `f(u, h, p..., t)` and hence `p_start` would need to be `3`.
+
+The returned function is `identity` if the system does not have an `IndexCache`.
+"""
+function wrap_mtkparameters(sys::AbstractSystem, isscalar::Bool, p_start = 2)
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
         offset = Int(is_time_dependent(sys))
 
         if isscalar
             function (expr)
-                p = gensym(:p)
+                destructured_args = [x for x in expr.args[p_start:end-offset] if x isa DestructuredArgs]
                 Func(
                     [
-                        expr.args[1],
+                        expr.args[begin:p_start-1]...,
                         DestructuredArgs(
-                            [arg.name for arg in expr.args[2:(end - offset)]], p),
-                        (isone(offset) ? (expr.args[end],) : ())...
+                            [arg.name for arg in expr.args[p_start:(end - offset)]], MTKPARAMETERS_ARG),
+                        expr.args[end-offset+1:end]...
                     ],
                     [],
-                    Let(expr.args[2:(end - offset)], expr.body, false)
+                    Let(destructured_args, expr.body, false)
                 )
             end
         else
             function (expr)
-                p = gensym(:p)
+                destructured_args = [x for x in expr.args[p_start:end-offset] if x isa DestructuredArgs]
                 Func(
                     [
-                        expr.args[1],
+                        expr.args[begin:p_start-1]...,
                         DestructuredArgs(
-                            [arg.name for arg in expr.args[2:(end - offset)]], p),
-                        (isone(offset) ? (expr.args[end],) : ())...
+                            [arg.name for arg in expr.args[p_start:(end - offset)]], MTKPARAMETERS_ARG),
+                        expr.args[end-offset+1:end]...
                     ],
                     [],
-                    Let(expr.args[2:(end - offset)], expr.body, false)
+                    Let(destructured_args, expr.body, false)
                 )
             end,
             function (expr)
-                p = gensym(:p)
+                destructured_args = [x for x in expr.args[p_start:end-offset] if x isa DestructuredArgs]
                 Func(
                     [
-                        expr.args[1],
-                        expr.args[2],
+                        expr.args[begin:p_start]...,
                         DestructuredArgs(
-                            [arg.name for arg in expr.args[3:(end - offset)]], p),
-                        (isone(offset) ? (expr.args[end],) : ())...
+                            [arg.name for arg in expr.args[p_start+1:(end - offset)]], MTKPARAMETERS_ARG),
+                        expr.args[end-offset+1:end]...
                     ],
                     [],
-                    Let(expr.args[3:(end - offset)], expr.body, false)
+                    Let(destructured_args, expr.body, false)
                 )
             end
         end