handle function wrapping no-recompile

SciML · Aug 23, 2022 · c1f1b27 · c1f1b27
1 parent 3d7c9b6
commit c1f1b27
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "JumpProcesses"
 uuid = "ccbc3e58-028d-4f4c-8cd5-9ae44345cda5"
 authors = ["Chris Rackauckas <[email protected]>"]
-version = "9.1.0"
+version = "9.1.1"
 
 [deps]
 ArrayInterfaceCore = "30b0a656-2188-435a-8636-2ec0e6a096e2"

diff --git a/src/problem.jl b/src/problem.jl
@@ -227,9 +227,17 @@ function extend_problem(prob::DiffEqBase.AbstractDiscreteProblem, jumps; rng = D
 end
 
 function extend_problem(prob::DiffEqBase.AbstractODEProblem, jumps; rng = DEFAULT_RNG)
-    function jump_f(du::ExtendedJumpArray, u::ExtendedJumpArray, p, t)
-        prob.f(du.u, u.u, p, t)
-        update_jumps!(du, u, p, t, length(u.u), jumps.variable_jumps...)
+    if isdefined(SciMLBase,:unwrapped_f)
+        _f = SciMLBase.unwrapped_f(prob.f)
+    else
+        _f = prob.f
+    end
+
+    jump_f = let _f = _f
+        function jump_f(du::ExtendedJumpArray, u::ExtendedJumpArray, p, t)
+            _f(du.u, u.u, p, t)
+            update_jumps!(du, u, p, t, length(u.u), jumps.variable_jumps...)
+        end
     end
     ttype = eltype(prob.tspan)
     u0 = ExtendedJumpArray(prob.u0,