Introduce ModelFunction type to dispatch on

Author: phipsgabler

src/compiler.jl

@@ -206,7 +206,7 @@ function build_model_info(input_expr)
             :vi => gensym(:vi),
             :sampler => gensym(:sampler),
             :model => gensym(:model),
-            :inner_function => gensym(:inner_function),
+            :model_function => gensym(modeldef[:name]),
             :defaults => gensym(:defaults)
         )
     )
@@ -319,6 +319,9 @@ function replace_tilde!(model_info)
 end
 """ |> Meta.parse |> eval
 
+# """ Unbreak code highlighting in Emacs julia-mode
+
+
 """
     generate_tilde(left, right, model_info)
 
@@ -437,7 +440,7 @@ function build_output(model_info)
     vi = main_body_names[:vi]
     model = main_body_names[:model]
     sampler = main_body_names[:sampler]
-    inner_function = main_body_names[:inner_function]
+    model_function = main_body_names[:model_function]
 
     # Arguments with default values
     args = model_info[:args]
@@ -481,28 +484,32 @@ function build_output(model_info)
     end
 
     ex = quote
+        function (::DynamicPPL.ModelFunction{$(QuoteNode(model_function))})(
+            $vi::DynamicPPL.VarInfo,
+            $sampler::DynamicPPL.AbstractSampler,
+            $ctx::DynamicPPL.AbstractContext,
+            $model
+        )
+            $unwrap_data_expr
+            DynamicPPL.resetlogp!($vi)
+            $main_body
+        end
+        
+        $model_function = DynamicPPL.ModelFunction{$(QuoteNode(model_function))}()
+        
         function $outer_function($(args...))
-            function $inner_function(
-                $vi::DynamicPPL.VarInfo,
-                $sampler::DynamicPPL.AbstractSampler,
-                $ctx::DynamicPPL.AbstractContext,
-                $model
-            )
-                $unwrap_data_expr
-                DynamicPPL.resetlogp!($vi)
-                $main_body
-            end
-            return DynamicPPL.Model($inner_function, $args_nt, $model_gen_constructor)
+            return DynamicPPL.Model($model_function, $args_nt, $model_gen_constructor)
         end
+        
         $model_gen = $model_gen_constructor
     end
 
     if !isempty(args)
-        ex = quote
-            $ex
-            # Allows passing arguments as kwargs
+        # Allows passing arguments as kwargs
+        kwform = quote
             $outer_function(;$(args...)) = $outer_function($(arg_syms...))
         end
+        push!(ex.args, kwform)
     end
 
     return esc(ex)

src/model.jl

@@ -1,6 +1,10 @@
+"""Tagged type for implementation of inner model functions that allow dispatch on their type."""
+struct ModelFunction{S} end
+
+
 """
-    struct Model{F, Targs <: NamedTuple, Tmodelgen, Tmissings <: Val}
-        f::F
+    struct Model{S, Targs <: NamedTuple, Tmodelgen, Tmissings <: Val}
+        f::ModelFunction{S}
         args::Targs
         modelgen::Tmodelgen
         missings::Tmissings
@@ -12,13 +16,13 @@ argument in `args` with a value `missing` will be in `missings` by default. Howe
 non-traditional use-cases `missings` can be defined differently. All variables in
 `missings` are treated as random variables rather than observations.
 """
-struct Model{F, Targs <: NamedTuple, Tmodelgen, Tmissings <: Val} <: AbstractModel
-    f::F
+struct Model{S, Targs <: NamedTuple, Tmodelgen, Tmissings <: Val} <: AbstractModel
+    f::ModelFunction{S}
     args::Targs
     modelgen::Tmodelgen
     missings::Tmissings
 end
-Model(f, args::NamedTuple, modelgen) = Model(f, args, modelgen, getmissing(args))
+Model(f::ModelFunction, args::NamedTuple, modelgen) = Model(f, args, modelgen, getmissing(args))
 (model::Model)(vi) = model(vi, SampleFromPrior())
 (model::Model)(vi, spl) = model(vi, spl, DefaultContext())
 (model::Model)(args...; kwargs...) = model.f(args..., model; kwargs...)