Merge pull request #445 from dpad/track-external-registrations

Track external registrations of Operations, so they can be injected into generated functions.

Author: ChrisRackauckas

src/build_function.jl

@@ -127,10 +127,26 @@ function _build_function(target::JuliaTarget, op::Operation, args...;
     if expression == Val{true}
         return oop_ex
     else
-        return GeneralizedGenerated.mk_function(@__MODULE__,oop_ex)
+        _build_and_inject_function(@__MODULE__, oop_ex)
     end
 end
 
+function _build_and_inject_function(mod::Module, ex)
+    # Generate the function, which will process the expression
+    runtimefn = GeneralizedGenerated.mk_function(@__MODULE__, ex)
+
+    # Extract the processed expression of the function body
+    params = typeof(runtimefn).parameters
+    fn_expr = GeneralizedGenerated.NGG.from_type(params[3])
+
+    # Inject our externally registered module functions 
+    new_expr = ModelingToolkit.inject_registered_module_functions(fn_expr)
+
+    # Reconstruct the RuntimeFn's Body
+    new_body = GeneralizedGenerated.NGG.to_type(new_expr)
+    return GeneralizedGenerated.RuntimeFn{params[1:2]..., new_body, params[4]}()
+end
+
 # Detect heterogeneous element types of "arrays of matrices/sparce matrices"
 function is_array_matrix(F)
     return isa(F, AbstractVector) && all(x->isa(x, AbstractArray), F)
@@ -288,7 +304,7 @@ function _build_function(target::JuliaTarget, rhss, args...;
     if expression == Val{true}
         return oop_ex, iip_ex
     else
-        return GeneralizedGenerated.mk_function(@__MODULE__,oop_ex), GeneralizedGenerated.mk_function(@__MODULE__,iip_ex)
+        return _build_and_inject_function(@__MODULE__, oop_ex), _build_and_inject_function(@__MODULE__, iip_ex)
     end
 end
 

src/function_registration.jl

@@ -14,12 +14,30 @@ registers `f` as a possible two-argument function.
 macro register(sig)
     splitsig = splitdef(:($sig = nothing))
     name = splitsig[:name]
+
+    # Extract the module and function name from the signature
+    if name isa Symbol
+        mod = __module__  # Calling module
+        funcname = name
+    else
+        mod = name.args[1]
+        funcname = name.args[2].value
+    end
+
     args = splitsig[:args]
     typargs = typed_args(args)
     defs = :()
     for typarg in typargs
         splitsig[:args] = typarg
-        splitsig[:body] = :(Operation($name, Expression[$(args...)]))
+        if mod == (@__MODULE__)  # If the calling module is ModelingToolkit itself...
+            splitsig[:body] = :(Operation($name, Expression[$(args...)]))
+        else
+            # Register the function's associated model so we can inject it in later.
+            splitsig[:body] = quote
+                get!(ModelingToolkit.registered_external_functions, Symbol($("$funcname")), $mod)
+                Operation($name, Expression[$(args...)])
+            end
+        end
         defs = :($defs; $(combinedef(splitsig)))
     end
     esc(defs)
@@ -69,3 +87,30 @@ Base.:^(x::Expression,y::T) where T <: Rational = Operation(Base.:^, Expression[
 
 Base.getindex(x::Operation,i::Int64) = Operation(getindex,[x,i])
 Base.one(::Operation) = 1
+
+# Ensure that Operations that get @registered from outside the ModelingToolkit
+# module can work without having to bring in the associated function into the
+# ModelingToolkit namespace. We basically store information about functions
+# registered at runtime in a ModelingToolkit variable,
+# `registered_external_functions`. It's not pretty, but we are limited by the
+# way GeneralizedGenerated builds a function (adding "ModelingToolkit" to every
+# function call).
+# ---
+const registered_external_functions = Dict{Symbol,Module}()
+function inject_registered_module_functions(expr)
+    MacroTools.postwalk(expr) do x
+        # We need to find all function calls in the expression...
+        MacroTools.@capture(x, f_(xs__))  
+
+        if !isnothing(f) && f isa Expr && f.head == :. && f.args[2] isa QuoteNode
+            # If the function call matches any of the functions we've
+            # registered, set the calling module (which is probably
+            # "ModelingToolkit") to the module it is registered to.
+            f_name = f.args[2].value  # function name
+            f.args[1] = get(registered_external_functions, f_name, f.args[1])
+        end
+
+        # Make sure we rebuild the expression as is.
+        return x
+    end
+end

test/function_registration.jl

@@ -0,0 +1,59 @@
+# TEST: Function registration in a module.
+# ------------------------------------------------
+module MyModule
+    using ModelingToolkit, DiffEqBase, LinearAlgebra, Test
+    @parameters t x
+    @variables u(t)
+    @derivatives Dt'~t
+
+    function do_something(a)
+        a + 10
+    end
+    @register do_something(a)
+
+    eq  = Dt(u) ~ do_something(x)
+    sys = ODESystem([eq], t, [u], [x])
+    fun = ODEFunction(sys)
+
+    @test fun([0.5], [5.0], 0.) == [15.0]
+end
+
+# TEST: Function registration in a nested module.
+# ------------------------------------------------
+module MyModule2
+    module MyNestedModule
+        using ModelingToolkit, DiffEqBase, LinearAlgebra, Test
+        @parameters t x
+        @variables u(t)
+        @derivatives Dt'~t
+
+        function do_something_2(a)
+            a + 20
+        end
+        @register do_something_2(a)
+
+        eq  = Dt(u) ~ do_something_2(x)
+        sys = ODESystem([eq], t, [u], [x])
+        fun = ODEFunction(sys)
+
+        @test fun([0.5], [3.0], 0.) == [23.0]
+    end
+end
+
+# TEST: Function registration outside any modules.
+# ------------------------------------------------
+using ModelingToolkit, DiffEqBase, LinearAlgebra, Test
+@parameters t x
+@variables u(t)
+@derivatives Dt'~t
+
+function do_something_3(a)
+    a + 30
+end
+@register do_something_3(a)
+
+eq  = Dt(u) ~ do_something_3(x)
+sys = ODESystem([eq], t, [u], [x])
+fun = ODEFunction(sys)
+
+@test fun([0.5], [7.0], 0.) == [37.0]

test/runtests.jl

@@ -20,6 +20,7 @@ using SafeTestsets, Test
 @safetestset "Lowering Integration Test" begin include("lowering_solving.jl") end
 @safetestset "Test Big System Usage" begin include("bigsystem.jl") end
 @safetestset "Depdendency Graph Test" begin include("dep_graphs.jl") end
+@safetestset "Function Registration Test" begin include("function_registration.jl") end
 #@testset "Latexify recipes Test" begin include("latexify.jl") end
 @testset "Distributed Test" begin include("distributed.jl") end
 @testset "Array of Array Test" begin include("build_function_arrayofarray.jl") end