Track external registrations of Operations, so they can be injected into generated functions.This solves the issue where an external package may wish to use a function it defines "my_func" as an Operation, and so calls "@register my_func(x)" on it. However, when ModelingToolkit eventually generates a function including that Operation, GeneralizedGenerated renames all calls to the function to "(ModelingToolkit).my_func", which does not exist because that function was not defined in the ModelingToolkit namespace. Instead, we track all the functions that are registered and which module they come from, and use that to re-write the generated functions with the appropriate module.

Author: dpad

Project.toml

@@ -12,6 +12,7 @@ DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 GeneralizedGenerated = "6b9d7cbe-bcb9-11e9-073f-15a7a543e2eb"
+JuliaVariables = "b14d175d-62b4-44ba-8fb7-3064adc8c3ec"
 Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
 LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -35,6 +36,7 @@ DiffEqJump = "6.7.5"
 DiffRules = "0.1, 1.0"
 DocStringExtensions = "0.7, 0.8"
 GeneralizedGenerated = "0.1.4, 0.2"
+JuliaVariables = "0.2.0"
 Latexify = "0.11, 0.12, 0.13"
 LightGraphs = "1.3"
 MacroTools = "0.5"

src/function_registration.jl

@@ -10,6 +10,12 @@ ModelingToolkit IR. Example:
 ```
 
 registers `f` as a possible two-argument function.
+
+NOTE: If registering outside of ModelingToolkit (i.e. in your own package),
+this should be done at runtime (e.g. in a package `__init__()`, or inside a
+method that is called at runtime and not during precompile) to ensure that
+any generated functions will use your registered method. See
+`inject_registered_module_functions()`.
 """
 macro register(sig)
     splitsig = splitdef(:($sig = nothing))
@@ -22,6 +28,11 @@ macro register(sig)
         splitsig[:body] = :(Operation($name, Expression[$(args...)]))
         defs = :($defs; $(combinedef(splitsig)))
     end
+    if (@__MODULE__) != ModelingToolkit
+        # Register external module registrations so we can rewrite any generated
+        # functions through JuliaVariables.solve().
+        get!(ModelingToolkit.registered_external_functions, name, @__MODULE__)
+    end
     esc(defs)
 end
 # Create all valid combinations of Expression,Number for function signature
@@ -69,3 +80,43 @@ 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).
+# ---
+import JuliaVariables
+const registered_external_functions = Dict{Symbol,Module}()
+function inject_registered_module_functions(expr)
+    MacroTools.postwalk(expr) do x
+        MacroTools.@capture(x, f_(xs__))  # We need to find all function calls in the expression.
+        # If the function call has been converted to a JuliaVariables.Var and matches
+        # one of the functions we've registered...
+        if !isnothing(f) && x.args[1] isa JuliaVariables.Var && x.args[1].name in keys(registered_external_functions)
+            # Rewrite it from a Var to a regular function call.
+            x.args[1] = getproperty(registered_external_functions[x.args[1].name], x.args[1].name)
+        end
+        return x  # Make sure we rebuild the expression as is.
+    end
+end
+
+# TODO: Overwriting this function works, but is quite ugly. Is there a nicer way to inject the module names?
+function GeneralizedGenerated.mk_function(mod::Module, ex)
+    ex = macroexpand(mod, ex)
+    ex = GeneralizedGenerated.simplify_ex(ex)
+    ex = GeneralizedGenerated.solve(ex)
+
+    # We need to modify the expression built by the JuliaVariables.solve(ex)
+    # method, before GeneralizedGenerated.closure_conv(mod, ex) converts it to a
+    # RuntimeFn (as done in GeneralizedGenerated.mk_function(mod, ex)).
+    ex = inject_registered_module_functions(ex)
+
+    fn = GeneralizedGenerated.closure_conv(mod, ex)
+    if !(fn isa GeneralizedGenerated.RuntimeFn)
+        error("Expect an unnamed function expression. ")
+    end
+    fn
+end