fix: use GeneratedFunctionWrapper in codegen

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -422,28 +422,7 @@ function SymbolicIndexingInterface.timeseries_parameter_index(sys::AbstractSyste
 end
 
 function SymbolicIndexingInterface.parameter_observed(sys::AbstractSystem, sym)
-    if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
-        rawobs = build_explicit_observed_function(
-            sys, sym; param_only = true, return_inplace = true)
-        if rawobs isa Tuple
-            if is_time_dependent(sys)
-                obsfn = let oop = rawobs[1], iip = rawobs[2]
-                    f1a(p, t) = oop(p, t)
-                    f1a(out, p, t) = iip(out, p, t)
-                end
-            else
-                obsfn = let oop = rawobs[1], iip = rawobs[2]
-                    f1b(p) = oop(p)
-                    f1b(out, p) = iip(out, p)
-                end
-            end
-        else
-            obsfn = rawobs
-        end
-    else
-        obsfn = build_explicit_observed_function(sys, sym; param_only = true)
-    end
-    return obsfn
+    return build_explicit_observed_function(sys, sym; param_only = true)
 end
 
 function has_observed_with_lhs(sys, sym)
@@ -579,18 +558,8 @@ function SymbolicIndexingInterface.observed(
             end
         end
     end
-    _fn = build_explicit_observed_function(
+    return build_explicit_observed_function(
         sys, sym; eval_expression, eval_module, checkbounds)
-
-    if is_time_dependent(sys)
-        return _fn
-    else
-        return let _fn = _fn
-            fn2(u, p) = _fn(u, p)
-            fn2(::Nothing, p) = _fn([], p)
-            fn2
-        end
-    end
 end
 
 function SymbolicIndexingInterface.default_values(sys::AbstractSystem)

src/systems/diffeqs/abstractodesystem.jl

@@ -320,9 +320,7 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         expression_module = eval_module, checkbounds = checkbounds,
         kwargs...)
     f_oop, f_iip = eval_or_rgf.(f_gen; eval_expression, eval_module)
-
-    f(u, p, t) = f_oop(u, p, t)
-    f(du, u, p, t) = f_iip(du, u, p, t)
+    f = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(f_oop, f_iip)
 
     if specialize === SciMLBase.FunctionWrapperSpecialize && iip
         if u0 === nothing || p === nothing || t === nothing
@@ -338,10 +336,7 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
             expression_module = eval_module,
             checkbounds = checkbounds, kwargs...)
         tgrad_oop, tgrad_iip = eval_or_rgf.(tgrad_gen; eval_expression, eval_module)
-
-        ___tgrad(u, p, t) = tgrad_oop(u, p, t)
-        ___tgrad(J, u, p, t) = tgrad_iip(J, u, p, t)
-        _tgrad = ___tgrad
+        _tgrad = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(tgrad_oop, tgrad_iip)
     else
         _tgrad = nothing
     end
@@ -354,8 +349,7 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
             checkbounds = checkbounds, kwargs...)
         jac_oop, jac_iip = eval_or_rgf.(jac_gen; eval_expression, eval_module)
 
-        _jac(u, p, t) = jac_oop(u, p, t)
-        _jac(J, u, p, t) = jac_iip(J, u, p, t)
+        _jac = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(jac_oop, jac_iip)
     else
         _jac = nothing
     end
@@ -435,8 +429,7 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
         expression_module = eval_module, checkbounds = checkbounds,
         kwargs...)
     f_oop, f_iip = eval_or_rgf.(f_gen; eval_expression, eval_module)
-    f(du, u, p, t) = f_oop(du, u, p, t)
-    f(out, du, u, p, t) = f_iip(out, du, u, p, t)
+    f = GeneratedFunctionWrapper{(3, 4, is_split(sys))}(f_oop, f_iip)
 
     if jac
         jac_gen = generate_dae_jacobian(sys, dvs, ps;
@@ -446,8 +439,7 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
             checkbounds = checkbounds, kwargs...)
         jac_oop, jac_iip = eval_or_rgf.(jac_gen; eval_expression, eval_module)
 
-        _jac(du, u, p, ˍ₋gamma, t) = jac_oop(du, u, p, ˍ₋gamma, t)
-        _jac(J, du, u, p, ˍ₋gamma, t) = jac_iip(J, du, u, p, ˍ₋gamma, t)
+        _jac = GeneratedFunctionWrapper{(3, 5, is_split(sys))}(jac_oop, jac_iip)
     else
         _jac = nothing
     end
@@ -496,8 +488,7 @@ function DiffEqBase.DDEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
         expression_module = eval_module, checkbounds = checkbounds,
         kwargs...)
     f_oop, f_iip = eval_or_rgf.(f_gen; eval_expression, eval_module)
-    f(u, h, p, t) = f_oop(u, h, p, t)
-    f(du, u, h, p, t) = f_iip(du, u, h, p, t)
+    f = GeneratedFunctionWrapper{(3, 4, is_split(sys))}(f_oop, f_iip)
 
     DDEFunction{iip}(f; sys = sys, initialization_data)
 end
@@ -521,14 +512,12 @@ function DiffEqBase.SDDEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys
         expression_module = eval_module, checkbounds = checkbounds,
         kwargs...)
     f_oop, f_iip = eval_or_rgf.(f_gen; eval_expression, eval_module)
-    f(u, h, p, t) = f_oop(u, h, p, t)
-    f(du, u, h, p, t) = f_iip(du, u, h, p, t)
+    f = GeneratedFunctionWrapper{(3, 4, is_split(sys))}(f_oop, f_iip)
 
     g_gen = generate_diffusion_function(sys, dvs, ps; expression = Val{true},
         isdde = true, kwargs...)
     g_oop, g_iip = eval_or_rgf.(g_gen; eval_expression, eval_module)
-    g(u, h, p, t) = g_oop(u, h, p, t)
-    g(du, u, h, p, t) = g_iip(du, u, h, p, t)
+    g = GeneratedFunctionWrapper{(3, 4, is_split(sys))}(g_oop, g_iip)
 
     SDDEFunction{iip}(f, g; sys = sys, initialization_data)
 end
@@ -549,13 +538,6 @@ variable and parameter vectors, respectively.
 """
 struct ODEFunctionExpr{iip, specialize} end
 
-struct ODEFunctionClosure{O, I} <: Function
-    f_oop::O
-    f_iip::I
-end
-(f::ODEFunctionClosure)(u, p, t) = f.f_oop(u, p, t)
-(f::ODEFunctionClosure)(du, u, p, t) = f.f_iip(du, u, p, t)
-
 function ODEFunctionExpr{iip, specialize}(sys::AbstractODESystem, dvs = unknowns(sys),
         ps = parameters(sys), u0 = nothing;
         version = nothing, tgrad = false,
@@ -572,13 +554,14 @@ function ODEFunctionExpr{iip, specialize}(sys::AbstractODESystem, dvs = unknowns
     f_oop, f_iip = generate_function(sys, dvs, ps; expression = Val{true}, kwargs...)
 
     fsym = gensym(:f)
-    _f = :($fsym = $ODEFunctionClosure($f_oop, $f_iip))
+    _f = :($fsym = $(GeneratedFunctionWrapper{(2, 3, is_split(sys))})($f_oop, $f_iip))
     tgradsym = gensym(:tgrad)
     if tgrad
         tgrad_oop, tgrad_iip = generate_tgrad(sys, dvs, ps;
             simplify = simplify,
             expression = Val{true}, kwargs...)
-        _tgrad = :($tgradsym = $ODEFunctionClosure($tgrad_oop, $tgrad_iip))
+        _tgrad = :($tgradsym = $(GeneratedFunctionWrapper{(2, 3, is_split(sys))})(
+            $tgrad_oop, $tgrad_iip))
     else
         _tgrad = :($tgradsym = nothing)
     end
@@ -588,7 +571,8 @@ function ODEFunctionExpr{iip, specialize}(sys::AbstractODESystem, dvs = unknowns
         jac_oop, jac_iip = generate_jacobian(sys, dvs, ps;
             sparse = sparse, simplify = simplify,
             expression = Val{true}, kwargs...)
-        _jac = :($jacsym = $ODEFunctionClosure($jac_oop, $jac_iip))
+        _jac = :($jacsym = $(GeneratedFunctionWrapper{(2, 3, is_split(sys))})(
+            $jac_oop, $jac_iip))
     else
         _jac = :($jacsym = nothing)
     end
@@ -647,13 +631,6 @@ variable and parameter vectors, respectively.
 """
 struct DAEFunctionExpr{iip} end
 
-struct DAEFunctionClosure{O, I} <: Function
-    f_oop::O
-    f_iip::I
-end
-(f::DAEFunctionClosure)(du, u, p, t) = f.f_oop(du, u, p, t)
-(f::DAEFunctionClosure)(out, du, u, p, t) = f.f_iip(out, du, u, p, t)
-
 function DAEFunctionExpr{iip}(sys::AbstractODESystem, dvs = unknowns(sys),
         ps = parameters(sys), u0 = nothing;
         version = nothing, tgrad = false,
@@ -667,7 +644,7 @@ function DAEFunctionExpr{iip}(sys::AbstractODESystem, dvs = unknowns(sys),
     f_oop, f_iip = generate_function(sys, dvs, ps; expression = Val{true},
         implicit_dae = true, kwargs...)
     fsym = gensym(:f)
-    _f = :($fsym = $DAEFunctionClosure($f_oop, $f_iip))
+    _f = :($fsym = $(GeneratedFunctionWrapper{(3, 4, is_split(sys))})($f_oop, $f_iip))
     ex = quote
         $_f
         ODEFunction{$iip}($fsym)
@@ -708,6 +685,7 @@ function SymbolicTstops(
         expression = Val{true},
         p_start = 1, p_end = length(rps), add_observed = false, force_SA = true)
     tstops = eval_or_rgf(tstops; eval_expression, eval_module)
+    tstops = GeneratedFunctionWrapper{(1, 3, is_split(sys))}(tstops, nothing)
     return SymbolicTstops(tstops)
 end
 

src/systems/diffeqs/odesystem.jl

@@ -521,9 +521,14 @@ function build_explicit_observed_function(sys, ts;
         output_type, mkarray, try_namespaced = true, expression = Val{true})
     if fns isa Tuple
         oop, iip = eval_or_rgf.(fns; eval_expression, eval_module)
-        return return_inplace ? (oop, iip) : oop
+        f = GeneratedFunctionWrapper{(
+            p_start, length(args) - length(ps) + 1, is_split(sys))}(oop, iip)
+        return return_inplace ? (f, f) : f
     else
-        return eval_or_rgf(fns; eval_expression, eval_module)
+        f = eval_or_rgf(fns; eval_expression, eval_module)
+        f = GeneratedFunctionWrapper{(
+            p_start, length(args) - length(ps) + 1, is_split(sys))}(f, nothing)
+        return f
     end
 end
 

src/systems/diffeqs/sdesystem.jl

@@ -604,18 +604,14 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
         kwargs...)
     g_oop, g_iip = eval_or_rgf.(g_gen; eval_expression, eval_module)
 
-    f(u, p, t) = f_oop(u, p, t)
-    f(du, u, p, t) = f_iip(du, u, p, t)
-    g(u, p, t) = g_oop(u, p, t)
-    g(du, u, p, t) = g_iip(du, u, p, t)
+    f = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(f_oop, f_iip)
+    g = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(g_oop, g_iip)
 
     if tgrad
         tgrad_gen = generate_tgrad(sys, dvs, ps; expression = Val{true},
             kwargs...)
         tgrad_oop, tgrad_iip = eval_or_rgf.(tgrad_gen; eval_expression, eval_module)
-
-        _tgrad(u, p, t) = tgrad_oop(u, p, t)
-        _tgrad(J, u, p, t) = tgrad_iip(J, u, p, t)
+        _tgrad = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(tgrad_oop, tgrad_iip)
     else
         _tgrad = nothing
     end
@@ -625,8 +621,7 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
             sparse = sparse, kwargs...)
         jac_oop, jac_iip = eval_or_rgf.(jac_gen; eval_expression, eval_module)
 
-        _jac(u, p, t) = jac_oop(u, p, t)
-        _jac(J, u, p, t) = jac_iip(J, u, p, t)
+        _jac = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(jac_oop, jac_iip)
     else
         _jac = nothing
     end
@@ -637,10 +632,8 @@ function DiffEqBase.SDEFunction{iip, specialize}(sys::SDESystem, dvs = unknowns(
         Wfact_oop, Wfact_iip = eval_or_rgf.(tmp_Wfact; eval_expression, eval_module)
         Wfact_oop_t, Wfact_iip_t = eval_or_rgf.(tmp_Wfact_t; eval_expression, eval_module)
 
-        _Wfact(u, p, dtgamma, t) = Wfact_oop(u, p, dtgamma, t)
-        _Wfact(W, u, p, dtgamma, t) = Wfact_iip(W, u, p, dtgamma, t)
-        _Wfact_t(u, p, dtgamma, t) = Wfact_oop_t(u, p, dtgamma, t)
-        _Wfact_t(W, u, p, dtgamma, t) = Wfact_iip_t(W, u, p, dtgamma, t)
+        _Wfact = GeneratedFunctionWrapper{(2, 4, is_split(sys))}(Wfact_oop, Wfact_iip)
+        _Wfact_t = GeneratedFunctionWrapper{(2, 4, is_split(sys))}(Wfact_oop_t, Wfact_iip_t)
     else
         _Wfact, _Wfact_t = nothing, nothing
     end

src/systems/discrete_system/discrete_system.jl

@@ -354,8 +354,7 @@ function SciMLBase.DiscreteFunction{iip, specialize}(
     f_gen = generate_function(sys, dvs, ps; expression = Val{true},
         expression_module = eval_module, kwargs...)
     f_oop, f_iip = eval_or_rgf.(f_gen; eval_expression, eval_module)
-    f(u, p, t) = f_oop(u, p, t)
-    f(du, u, p, t) = f_iip(du, u, p, t)
+    f = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(f_oop, f_iip)
 
     if specialize === SciMLBase.FunctionWrapperSpecialize && iip
         if u0 === nothing || p === nothing || t === nothing

src/systems/jumps/jumpsystem.jl

@@ -283,7 +283,7 @@ function generate_rate_function(js::JumpSystem, rate)
         rate = substitute(rate, csubs)
     end
     p = reorder_parameters(js)
-    rf = build_function_wrapper(js, rate, unknowns(js), p...,
+    build_function_wrapper(js, rate, unknowns(js), p...,
         get_iv(js),
         expression = Val{true})
 end
@@ -302,7 +302,7 @@ end
 function assemble_vrj(
         js, vrj, unknowntoid; eval_expression = false, eval_module = @__MODULE__)
     rate = eval_or_rgf(generate_rate_function(js, vrj.rate); eval_expression, eval_module)
-
+    rate = GeneratedFunctionWrapper{(2, 3, is_split(js))}(rate, nothing)
     outputvars = (value(affect.lhs) for affect in vrj.affect!)
     outputidxs = [unknowntoid[var] for var in outputvars]
     affect = eval_or_rgf(generate_affect_function(js, vrj.affect!, outputidxs);
@@ -326,7 +326,7 @@ end
 function assemble_crj(
         js, crj, unknowntoid; eval_expression = false, eval_module = @__MODULE__)
     rate = eval_or_rgf(generate_rate_function(js, crj.rate); eval_expression, eval_module)
-
+    rate = GeneratedFunctionWrapper{(2, 3, is_split(js))}(rate, nothing)
     outputvars = (value(affect.lhs) for affect in crj.affect!)
     outputidxs = [unknowntoid[var] for var in outputvars]
     affect = eval_or_rgf(generate_affect_function(js, crj.affect!, outputidxs);

src/systems/nonlinear/nonlinearsystem.jl

@@ -343,16 +343,14 @@ function SciMLBase.NonlinearFunction{iip}(sys::NonlinearSystem, dvs = unknowns(s
     end
     f_gen = generate_function(sys, dvs, ps; expression = Val{true}, kwargs...)
     f_oop, f_iip = eval_or_rgf.(f_gen; eval_expression, eval_module)
-    f(u, p) = f_oop(u, p)
-    f(du, u, p) = f_iip(du, u, p)
+    f = GeneratedFunctionWrapper{(2, 2, is_split(sys))}(f_oop, f_iip)
 
     if jac
         jac_gen = generate_jacobian(sys, dvs, ps;
             simplify = simplify, sparse = sparse,
             expression = Val{true}, kwargs...)
         jac_oop, jac_iip = eval_or_rgf.(jac_gen; eval_expression, eval_module)
-        _jac(u, p) = jac_oop(u, p)
-        _jac(J, u, p) = jac_iip(J, u, p)
+        _jac = GeneratedFunctionWrapper{(2, 2, is_split(sys))}(jac_oop, jac_iip)
     else
         _jac = nothing
     end
@@ -397,6 +395,7 @@ function SciMLBase.IntervalNonlinearFunction(
     f_gen = generate_function(
         sys, dvs, ps; expression = Val{true}, scalar = true, kwargs...)
     f = eval_or_rgf(f_gen; eval_expression, eval_module)
+    f = GeneratedFunctionWrapper{(2, 2, is_split(sys))}(f, nothing)
 
     observedfun = ObservedFunctionCache(
         sys; eval_expression, eval_module, checkbounds = get(kwargs, :checkbounds, false))
@@ -431,13 +430,14 @@ function NonlinearFunctionExpr{iip}(sys::NonlinearSystem, dvs = unknowns(sys),
     if !iscomplete(sys)
         error("A completed `NonlinearSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `NonlinearFunctionExpr`")
     end
-    idx = iip ? 2 : 1
-    f = generate_function(sys, dvs, ps; expression = Val{true}, kwargs...)[idx]
+    f_oop, f_iip = generate_function(sys, dvs, ps; expression = Val{true}, kwargs...)
+    f = :($(GeneratedFunctionWrapper{(2, 2, is_split(sys))})($f_oop, $f_iip))
 
     if jac
-        _jac = generate_jacobian(sys, dvs, ps;
+        jac_oop, jac_iip = generate_jacobian(sys, dvs, ps;
             sparse = sparse, simplify = simplify,
-            expression = Val{true}, kwargs...)[idx]
+            expression = Val{true}, kwargs...)
+        _jac = :($(GeneratedFunctionWrapper{(2, 2, is_split(sys))})($jac_oop, $jac_iip))
     else
         _jac = :nothing
     end
@@ -484,8 +484,7 @@ function IntervalNonlinearFunctionExpr(
     end
 
     f = generate_function(sys, dvs, ps; expression = Val{true}, scalar = true, kwargs...)
-
-    IntervalNonlinearFunction{false}(f; sys = sys)
+    f = :($(GeneratedFunctionWrapper{2, 2, is_split(sys)})($f, nothing))
 
     ex = quote
         f = $f
@@ -584,7 +583,9 @@ function CacheWriter(sys::AbstractSystem, buffer_types::Vector{TypeT},
         DestructuredArgs.(rps)...; p_start = 3, p_end = length(rps) + 2,
         expression = Val{true}, add_observed = false,
         extra_assignments = [obs_assigns; body])
-    return CacheWriter(eval_or_rgf(fn; eval_expression, eval_module))
+    fn = eval_or_rgf(fn; eval_expression, eval_module)
+    fn = GeneratedFunctionWrapper{(3, 3, is_split(sys))}(fn, nothing)
+    return CacheWriter(fn)
 end
 
 struct SCCNonlinearFunction{iip} end
@@ -603,9 +604,7 @@ function SCCNonlinearFunction{iip}(
         p_end = length(rps) + length(cachesyms) + 1, add_observed = false,
         extra_assignments = obs_assignments, expression = Val{true})
     f_oop, f_iip = eval_or_rgf.(f_gen; eval_expression, eval_module)
-
-    f(u, p) = f_oop(u, p)
-    f(resid, u, p) = f_iip(resid, u, p)
+    f = GeneratedFunctionWrapper{(2, 2, is_split(sys))}(f_oop, f_iip)
 
     subsys = NonlinearSystem(_eqs, _dvs, ps; observed = _obs,
         parameter_dependencies = parameter_dependencies(sys), name = nameof(sys))