Inject registered functions into expressions used to build EvalFunc. Fixes the regression of #445 caused by #451.

Author: dpad

src/build_function.jl

@@ -125,15 +125,15 @@ function _build_function(target::JuliaTarget, op::Operation, args...;
     end
 
     if expression == Val{true}
-        return oop_ex
+        return ModelingToolkit.inject_registered_module_functions(oop_ex)
     else
         _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)
+    runtimefn = GeneralizedGenerated.mk_function(mod, ex)
 
     # Extract the processed expression of the function body
     params = typeof(runtimefn).parameters
@@ -339,7 +339,7 @@ function _build_function(target::JuliaTarget, rhss, args...;
     end
 
     if expression == Val{true}
-        return oop_ex, iip_ex
+        return ModelingToolkit.inject_registered_module_functions(oop_ex), ModelingToolkit.inject_registered_module_functions(iip_ex)
     else
         return _build_and_inject_function(@__MODULE__, oop_ex), _build_and_inject_function(@__MODULE__, iip_ex)
     end

src/function_registration.jl

@@ -99,18 +99,21 @@ Base.one(::Operation) = 1
 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__))  
+        # Find all function calls in the expression and extract the function
+        # name and calling module.
+        MacroTools.@capture(x, f_module_.f_name_(xs__))
+        if isnothing(f_module)
+            MacroTools.@capture(x, f_name_(xs__))
+        end
 
-        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])
+        if !isnothing(f_name)
+            # Set the calling module to the module that registered it.
+            mod = get(registered_external_functions, f_name, f_module)
+            if !isnothing(mod)
+                x.args[1] = :(getproperty($mod, $(Meta.quot(f_name))))
+            end
         end
 
-        # Make sure we rebuild the expression as is.
         return x
     end
 end

test/function_registration.jl

@@ -1,3 +1,8 @@
+# Each test here builds an ODEFunction including some user-registered
+# Operations. The test simply checks that calling the ODEFunction
+# appropriately calls the registered functions, whether the call is
+# qualified (with a module name) or not.
+
 # TEST: Function registration in a module.
 # ------------------------------------------------
 module MyModule
@@ -11,11 +16,11 @@ module MyModule
     end
     @register do_something(a)
 
-    eq  = Dt(u) ~ do_something(x)
+    eq  = Dt(u) ~ do_something(x) + MyModule.do_something(x)
     sys = ODESystem([eq], t, [u], [x])
     fun = ODEFunction(sys)
 
-    @test_broken fun([0.5], [5.0], 0.) == [15.0]
+    @test fun([0.5], [5.0], 0.) == [30.0]
 end
 
 # TEST: Function registration in a nested module.
@@ -32,11 +37,11 @@ module MyModule2
         end
         @register do_something_2(a)
 
-        eq  = Dt(u) ~ do_something_2(x)
+        eq  = Dt(u) ~ do_something_2(x) + MyNestedModule.do_something_2(x)
         sys = ODESystem([eq], t, [u], [x])
         fun = ODEFunction(sys)
 
-        @test_broken fun([0.5], [3.0], 0.) == [23.0]
+        @test fun([0.5], [3.0], 0.) == [46.0]
     end
 end
 
@@ -52,8 +57,8 @@ function do_something_3(a)
 end
 @register do_something_3(a)
 
-eq  = Dt(u) ~ do_something_3(x)
+eq  = Dt(u) ~ do_something_3(x) + (@__MODULE__).do_something_3(x)
 sys = ODESystem([eq], t, [u], [x])
 fun = ODEFunction(sys)
 
-@test_broken fun([0.5], [7.0], 0.) == [37.0]
+@test fun([0.5], [7.0], 0.) == [74.0]