Wrap ODE functions in a fully transparent wrapper function

Author: oscarddssmith

src/problems/ode_problems.jl

@@ -198,7 +198,12 @@ function ODEProblem(f, u0, tspan, p = NullParameters(); kwargs...)
     iip = isinplace(f, 4)
     _u0 = prepare_initial_state(u0)
     _tspan = promote_tspan(tspan)
-    _f = ODEFunction{iip, DEFAULT_SPECIALIZATION}(f)
+    _f = if iip
+        out = copy(u0) # TODO: do this properly
+        ODEFunction{iip, FullSpecialize}(DUMB_WRAPPER(out, u0, p, first(_tspan), ODE_F_WRAPPER(f)))
+    else
+        _f = ODEFunction{iip, DEFAULT_SPECIALIZATION}(f)
+    end
     ODEProblem(_f, _u0, _tspan, p; kwargs...)
 end
 

src/scimlfunctions.jl

@@ -2415,8 +2415,31 @@ end
 (f::BVPFunction)(args...) = f.f(args...)
 (f::DynamicalBVPFunction)(args...) = f.f(args...)
 
-######### Basic Constructor
+#### fun hack to make NOSPECIALIZE fast
+mutable struct ODE_F_WRAPPER{F}
+    const f::F
+end
+mutable struct DUMB_WRAPPER{OUT, U, P, T}
+    out::OUT
+    u::U
+    p::P
+    t::T
+    f::ODE_F_WRAPPER
+end
+function (w::ODE_F_WRAPPER)(d::DUMB_WRAPPER)
+    w.f(d.out, d.u, d.p, d.t)
+    return nothing
+end
+function (w::DUMB_WRAPPER)(out, u, p, t)
+    w.out = out
+    w.u = u
+    w.p = p
+    w.t = t
+    w.f(w)
+    return nothing
+end
 
+######### Basic Constructor
 function ODEFunction{iip, specialize}(f;
         mass_matrix = __has_mass_matrix(f) ? f.mass_matrix :
                       I,