Merge branch 'master' into phg/split_tilde

Author: phipsgabler

src/DynamicPPL.jl

@@ -61,8 +61,12 @@ export  VarName,
         vectorize,
         set_resume!,
 # Model
+        ModelGen,
         Model,
-        getmissing,
+        getmissings,
+        getargnames,
+        getdefaults,
+        getgenerator,
         runmodel!,
 # Samplers
         Sampler,
@@ -91,6 +95,11 @@ const DEBUG = Bool(parse(Int, get(ENV, "DEBUG_DYNAMICPPL", "0")))
 # Used here and overloaded in Turing
 function getspace end
 
+# Necessary forward declarations
+abstract type AbstractVarInfo end
+abstract type AbstractContext end
+
+
 include("utils.jl")
 include("selector.jl")
 include("model.jl")

src/compiler.jl

@@ -1,25 +1,3 @@
-"""
-    struct ModelGen{Targs, F, Tdefaults} <: Function
-        f::F
-        defaults::Tdefaults
-    end
-
-A `Model` generator. This is the output of the `@model` macro. `Targs` is the tuple 
-of the symbols of the model's arguments. `defaults` is the `NamedTuple` of default values
-of the arguments, if any. Every `ModelGen` is callable with the arguments `Targs`, 
-returning an instance of `Model`.
-"""
-struct ModelGen{Targs, F, Tdefaults} <: Function
-    f::F
-    defaults::Tdefaults
-end
-ModelGen{Targs}(args...) where {Targs} = ModelGen{Targs, typeof.(args)...}(args...)
-(m::ModelGen)(args...; kwargs...) = m.f(args...; kwargs...)
-function Base.getproperty(m::ModelGen{Targs}, f::Symbol) where {Targs}
-    f === :args && return Targs
-    return Base.getfield(m, f)
-end
-
 macro varinfo()
     :(throw(_error_msg()))
 end
@@ -61,18 +39,18 @@ When `ex` is not a variable (e.g., a literal), the function returns `false` as w
 macro isassumption(data_vars, missing_vars, ex)
     :false
 end
-macro isassumption(data_vars, missing_vars, ex::Union{Symbol, Expr})
+macro isassumption(model, ex::Union{Symbol, Expr})
     sym = gensym(:sym)
     lhs = gensym(:lhs)
     return esc(quote
         # Extract symbol
         $sym = Val($(vsym(ex)))
         # This branch should compile nicely in all cases except for partial missing data
         # For example, when `ex` is `x[i]` and `x isa Vector{Union{Missing, Float64}}`
-        if !DynamicPPL.inparams($sym, $data_vars) || DynamicPPL.inparams($sym, $missing_vars)
+        if !DynamicPPL.inargnames($sym, $model) || DynamicPPL.inmissings($sym, $model)
             true
         else
-            if DynamicPPL.inparams($sym, $data_vars)
+            if DynamicPPL.inargnames($sym, $model)
                 # Evaluate the lhs
                 $lhs = $ex
                 if $lhs === missing
@@ -87,10 +65,6 @@ macro isassumption(data_vars, missing_vars, ex::Union{Symbol, Expr})
     end)
 end
 
-@generated function inparams(::Val{s}, ::Val{t}) where {s, t}
-    return (s in t) ? :(true) : :(false)
-end
-
 #################
 # Main Compiler #
 #################
@@ -152,7 +126,7 @@ function build_model_info(input_expr)
     else
         nt_type = Expr(:curly, :NamedTuple, 
             Expr(:tuple, QuoteNode.(arg_syms)...), 
-            Expr(:curly, :Tuple, [:(DynamicPPL.get_type($x)) for x in arg_syms]...)
+            Expr(:curly, :Tuple, [:(Core.Typeof($x)) for x in arg_syms]...)
         )
         args_nt = Expr(:call, :(DynamicPPL.namedtuple), nt_type, Expr(:tuple, arg_syms...))
     end
@@ -206,27 +180,13 @@ function build_model_info(input_expr)
             :ctx => gensym(:ctx),
             :vi => gensym(:vi),
             :sampler => gensym(:sampler),
-            :model => gensym(:model),
-            :inner_function => gensym(:inner_function),
-            :defaults => gensym(:defaults)
+            :model => gensym(:model)
         )
     )
 
     return model_info
 end
 
-function to_namedtuple_expr(syms::Vector, vals = syms)
-    if length(syms) == 0
-        nt = :(NamedTuple())
-    else
-        nt_type = Expr(:curly, :NamedTuple, 
-            Expr(:tuple, QuoteNode.(syms)...), 
-            Expr(:curly, :Tuple, [:(DynamicPPL.get_type($x)) for x in vals]...)
-        )
-        nt = Expr(:call, :(DynamicPPL.namedtuple), nt_type, Expr(:tuple, vals...))
-    end
-    return nt
-end
 
 """
     replace_vi!(model_info)
@@ -330,19 +290,16 @@ The `tilde` function generates `observe` expression for data variables and `assu
 expressions for parameter variables, updating `model_info` in the process.
 """
 function generate_tilde(left, right, model_info)
-    arg_syms = Val((model_info[:arg_syms]...,))
     model = model_info[:main_body_names][:model]
     vi = model_info[:main_body_names][:vi]
     ctx = model_info[:main_body_names][:ctx]
     sampler = model_info[:main_body_names][:sampler]
-
-    @gensym(out,
-            lp,
-            vn,
-            inds,
-            isassumption,
-            temp_right)
-    
+    temp_right = gensym(:temp_right)
+    out = gensym(:out)
+    lp = gensym(:lp)
+    vn = gensym(:vn)
+    inds = gensym(:inds)
+    isassumption = gensym(:isassumption)
     assert_ex = :(DynamicPPL.assert_dist($temp_right, msg = $(wrong_dist_errormsg(@__LINE__))))
 
     if left isa Symbol || left isa Expr
@@ -351,7 +308,7 @@ function generate_tilde(left, right, model_info)
             $assert_ex
 
             $vn, $inds = $(varname(left)), $(vinds(left))
-            $isassumption = DynamicPPL.@isassumption($arg_syms, DynamicPPL.getmissing($model), $left)
+            $isassumption = DynamicPPL.@isassumption($model, $left)
             if $isassumption 
                 $out = DynamicPPL.tilde_assume($ctx, $sampler, $temp_right, $vn, $inds, $vi)
                 $left = $out[1]
@@ -386,20 +343,16 @@ This function returns the expression that replaces `left .~ right` in the model
 will be run.
 """
 function generate_dot_tilde(left, right, model_info)
-    arg_syms = Val((model_info[:arg_syms]...,))
     model = model_info[:main_body_names][:model]
     vi = model_info[:main_body_names][:vi]
     ctx = model_info[:main_body_names][:ctx]
     sampler = model_info[:main_body_names][:sampler]
-    
-    @gensym(out,
-            preprocessed,
-            lp,
-            vn,
-            inds,
-            isassumption,
-            temp_right)
-    
+    out = gensym(:out)
+    temp_right = gensym(:temp_right)
+    isassumption = gensym(:isassumption)
+    lp = gensym(:lp)
+    vn = gensym(:vn)
+    inds = gensym(:inds)
     assert_ex = :(DynamicPPL.assert_dist($temp_right, msg = $(wrong_dist_errormsg(@__LINE__))))
 
     if left isa Symbol || left isa Expr
@@ -408,7 +361,7 @@ function generate_dot_tilde(left, right, model_info)
             $assert_ex
 
             $vn, $inds = $(varname(left)), $(vinds(left))
-            $isassumption = DynamicPPL.@isassumption($arg_syms, DynamicPPL.getmissing($model), $left)
+            $isassumption = DynamicPPL.@isassumption($model, $left)
 
             if $isassumption
                 $out = DynamicPPL.dot_tilde_assume($ctx, $sampler, $temp_right, $left, $vn, $inds, $vi)
@@ -453,7 +406,6 @@ 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]
 
     # Arguments with default values
     args = model_info[:args]
@@ -468,16 +420,9 @@ function build_output(model_info)
     whereparams = model_info[:whereparams]
     # Model generator name
     model_gen = model_info[:name]
-    # Outer function name
-    outer_function = gensym(model_info[:name])
     # Main body of the model
     main_body = model_info[:main_body]
-    model_gen_constructor = quote
-        DynamicPPL.ModelGen{$(Tuple(arg_syms))}(
-            $outer_function, 
-            $defaults_nt,
-        )
-    end
+
     unwrap_data_expr = Expr(:block)
     for var in arg_syms
         temp_var = gensym(:temp_var)
@@ -496,40 +441,32 @@ function build_output(model_info)
         end)
     end
 
+    @gensym(evaluator, generator)
+    generator_kw_form = isempty(args) ? () : (:($generator(;$(args...)) = $generator($(arg_syms...))),)
+    model_gen_constructor = :(DynamicPPL.ModelGen{$(Tuple(arg_syms))}($generator, $defaults_nt))
+    
     ex = quote
-        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)
+        function $evaluator(
+            $model::Model,
+            $vi::DynamicPPL.VarInfo,
+            $sampler::DynamicPPL.AbstractSampler,
+            $ctx::DynamicPPL.AbstractContext,
+        )
+            $unwrap_data_expr
+            DynamicPPL.resetlogp!($vi)
+            $main_body
         end
-        $model_gen = $model_gen_constructor
-    end
+        
 
-    if !isempty(args)
-        ex = quote
-            $ex
-            # Allows passing arguments as kwargs
-            $outer_function(;$(args...)) = $outer_function($(arg_syms...))
-        end
+        $generator($(args...)) = DynamicPPL.Model($evaluator, $args_nt, $model_gen_constructor)
+        $(generator_kw_form...)
+        
+        $model_gen = $model_gen_constructor
     end
 
     return esc(ex)
 end
 
-# A hack for NamedTuple type specialization
-# (T = Int,) has type NamedTuple{(:T,), Tuple{DataType}} by default
-# With this function, we can make it NamedTuple{(:T,), Tuple{Type{Int}}}
-# Both are correct, but the latter is what we want for type stability
-get_type(::Type{T}) where {T} = Type{T}
-get_type(t) = typeof(t)
 
 function warn_empty(body)
     if all(l -> isa(l, LineNumberNode), body.args)

src/contexts.jl

@@ -1,5 +1,3 @@
-abstract type AbstractContext end
-
 """
     struct DefaultContext <: AbstractContext end