Support for immutable AbstractVarInfo

Project.toml

@@ -11,6 +11,7 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
@@ -21,5 +22,6 @@ Bijectors = "0.5.2, 0.6, 0.7, 0.8, 0.9"
 ChainRulesCore = "0.9.7, 0.10"
 Distributions = "0.23.8, 0.24, 0.25"
 MacroTools = "0.5.6"
+Setfield = "0.7"
 ZygoteRules = "0.2"
 julia = "1.3"

src/DynamicPPL.jl

@@ -10,6 +10,7 @@ using ChainRulesCore: ChainRulesCore
 using MacroTools: MacroTools
 using ZygoteRules: ZygoteRules
 using BangBang: BangBang
+using Setfield: Setfield
 
 using Random: Random
 
@@ -28,15 +29,22 @@ import Base:
     keys,
     haskey
 
+import BangBang: push!!, empty!!, setindex!!
+
 # VarInfo
 export AbstractVarInfo,
     VarInfo,
     UntypedVarInfo,
     TypedVarInfo,
+    push!!,
+    empty!!,
     getlogp,
     setlogp!,
     acclogp!,
     resetlogp!,
+    setlogp!!,
+    acclogp!!,
+    resetlogp!!,
     get_num_produce,
     set_num_produce!,
     reset_num_produce!,
@@ -51,6 +59,8 @@ export AbstractVarInfo,
     istrans,
     link!,
     invlink!,
+    link!!,
+    invlink!!,
     tonamedtuple,
     # VarName (reexport from AbstractPPL)
     VarName,
@@ -134,4 +144,26 @@ include("compat/ad.jl")
 include("loglikelihoods.jl")
 include("submodel_macro.jl")
 
+# Deprecations
+@deprecate empty!(vi::VarInfo) empty!!(vi::VarInfo)
+@deprecate push!(vi::AbstractVarInfo, vn::VarName, r, dist::Distribution) push!!(
+    vi::AbstractVarInfo, vn::VarName, r, dist::Distribution
+)
+@deprecate push!(
+    vi::AbstractVarInfo, vn::VarName, r, dist::Distribution, sampler::AbstractSampler
+) push!!(vi::AbstractVarInfo, vn::VarName, r, dist::Distribution, sampler::AbstractSampler)
+@deprecate push!(vi::AbstractVarInfo, vn::VarName, r, dist::Distribution, gid::Selector) push!!(
+    vi::AbstractVarInfo, vn::VarName, r, dist::Distribution, gid::Selector
+)
+@deprecate push!(
+    vi::AbstractVarInfo, vn::VarName, r, dist::Distribution, gid::Set{Selector}
+) push!!(vi::AbstractVarInfo, vn::VarName, r, dist::Distribution, gid::Set{Selector})
+
+@deprecate setlogp!(vi, logp) setlogp!!(vi, logp)
+@deprecate acclogp!(vi, logp) acclogp!!(vi, logp)
+@deprecate resetlogp!(vi) resetlogp!!(vi)
+
+@deprecate link!(vi, spl) link!!(vi, spl)
+@deprecate invlink!(vi, spl) invlink!!(vi, spl)
+
 end # module

src/compat/ad.jl

@@ -1,5 +1,5 @@
 # See https://github.com/TuringLang/Turing.jl/issues/1199
-ChainRulesCore.@non_differentiable push!(
+ChainRulesCore.@non_differentiable push!!(
     vi::VarInfo, vn::VarName, r, dist::Distribution, gidset::Set{Selector}
 )
 

src/compiler.jl

@@ -369,7 +369,7 @@ function generate_tilde(left, right)
     # If the LHS is a literal, it is always an observation
     if isliteral(left)
         return quote
-            $(DynamicPPL.tilde_observe!)(
+            _, __varinfo__ = $(DynamicPPL.tilde_observe!!)(
                 __context__, $(DynamicPPL.check_tilde_rhs)($right), $left, __varinfo__
             )
         end
@@ -383,7 +383,7 @@ function generate_tilde(left, right)
         $inds = $(vinds(left))
         $isassumption = $(DynamicPPL.isassumption(left))
         if $isassumption
-            $left = $(DynamicPPL.tilde_assume!)(
+            $left, __varinfo__ = $(DynamicPPL.tilde_assume!!)(
                 __context__,
                 $(DynamicPPL.unwrap_right_vn)(
                     $(DynamicPPL.check_tilde_rhs)($right), $vn
@@ -397,7 +397,7 @@ function generate_tilde(left, right)
                 $left = $(DynamicPPL.getvalue_nested)(__context__, $vn)
             end
 
-            $(DynamicPPL.tilde_observe!)(
+            _, __varinfo__ = $(DynamicPPL.tilde_observe!!)(
                 __context__,
                 $(DynamicPPL.check_tilde_rhs)($right),
                 $(maybe_view(left)),
@@ -418,7 +418,7 @@ function generate_dot_tilde(left, right)
     # If the LHS is a literal, it is always an observation
     if isliteral(left)
         return quote
-            $(DynamicPPL.dot_tilde_observe!)(
+            _, __varinfo__ = $(DynamicPPL.dot_tilde_observe!!)(
                 __context__, $(DynamicPPL.check_tilde_rhs)($right), $left, __varinfo__
             )
         end
@@ -432,7 +432,7 @@ function generate_dot_tilde(left, right)
         $inds = $(vinds(left))
         $isassumption = $(DynamicPPL.isassumption(left))
         if $isassumption
-            $left .= $(DynamicPPL.dot_tilde_assume!)(
+            _, __varinfo__ = $(DynamicPPL.dot_tilde_assume!!)(
                 __context__,
                 $(DynamicPPL.unwrap_right_left_vns)(
                     $(DynamicPPL.check_tilde_rhs)($right), $(maybe_view(left)), $vn
@@ -446,7 +446,7 @@ function generate_dot_tilde(left, right)
                 $left .= $(DynamicPPL.getvalue_nested)(__context__, $vn)
             end
 
-            $(DynamicPPL.dot_tilde_observe!)(
+            _, __varinfo__ = $(DynamicPPL.dot_tilde_observe!!)(
                 __context__,
                 $(DynamicPPL.check_tilde_rhs)($right),
                 $(maybe_view(left)),
@@ -458,6 +458,108 @@ function generate_dot_tilde(left, right)
     end
 end
 
+"""
+    isfuncdef(expr)
+
+Return `true` if `expr` is any form of function definition, and `false` otherwise.
+"""
+function isfuncdef(e::Expr)
+    return if Meta.isexpr(e, :function)
+        # Classic `function f(...)`
+        true
+    elseif Meta.isexpr(e, :->)
+        # Anonymous functions/lambdas, e.g. `do` blocks or `->` defs.
+        true
+    elseif Meta.isexpr(e, :(=)) && Meta.isexpr(e.args[1], :call)
+        # Short function defs, e.g. `f(args...) = ...`.
+        true
+    else
+        false
+    end
+end
+
+"""
+    replace_returns(expr)
+
+Return `Expr` with all `return ...` statements replaced with
+`return ..., DynamicPPL.return_values(__varinfo__)`.
+
+Note that this method will _not_ replace `return` statements within function
+definitions. This is checked using [`isfuncdef`](@ref).
+"""
+replace_returns(e) = e
+replace_returns(e::Symbol) = e
+function replace_returns(e::Expr)
+    if isfuncdef(e)
+        return e
+    end
+
+    if Meta.isexpr(e, :return)
+        # NOTE: `return` always has an argument. In the case of
+        # `return`, the parsed expression will be `return nothing`.
+        # Hence we don't need any special handling for empty returns.
+        retval_expr = if length(e.args) > 1
+            Expr(:tuple, e.args...)
+        else
+            e.args[1]
+        end
+
+        return :(return $(DynamicPPL.return_values)($retval_expr, __varinfo__))
+    end
+
+    return Expr(e.head, map(replace_returns, e.args)...)
+end
+
+"""
+    return_values(retval, varinfo)
+
+Return `(retval, varinfo)` if `retval` is not a `Tuple` with second
+component being a `AbstractVarInfo`.
+
+Used together with [`replace_returns`](@ref), it handles the following case.
+
+# Example
+
+Suppose the following is the return-value:
+
+```julia
+return x ~ Normal()
+```
+
+Without `return_values`, once expanded in [`generated_mainbody!`](@ref), this would be
+
+```julia
+return (x, __varinfo__ = tilde_assume!!(...)), __varinfo__
+```
+
+i.e. the return-value of the model would end up `(x, __varinfo__), __varinfo__`
+which in turn would lead to a `(::Model)(args...)` call returning `(x, __varinfo__)`,
+breaking with the expectation of the user.
+
+In such a scenario `return_values` effectively results in the following
+
+```julia
+return x, __varinfo__ = tilde_assume!!(...)
+```
+
+preserving user expectation, as desired.
+"""
+return_values(retval, varinfo::AbstractVarInfo) = (retval, varinfo)
+return_values(retval::Tuple{<:Any,<:AbstractVarInfo}, ::AbstractVarInfo) = retval
+
+# If it's just a symbol, e.g. `f(x) = 1`, then we make it `f(x) = return 1`.
+make_returns_explicit!(body) = Expr(:return, body)
+function make_returns_explicit!(body::Expr)
+    # If the last statement is a return-statement, we don't do anything.
+    if Meta.isexpr(body.args[end], :return)
+        return body
+    end
+
+    # Otherwise we replace the last statement with a `return` statement.
+    body.args[end] = Expr(:return, body.args[end])
+    return body
+end
+
 const FloatOrArrayType = Type{<:Union{AbstractFloat,AbstractArray}}
 hasmissing(T::Type{<:AbstractArray{TA}}) where {TA<:AbstractArray} = hasmissing(TA)
 hasmissing(T::Type{<:AbstractArray{>:Missing}}) = true
@@ -489,7 +591,11 @@ function build_output(modelinfo, linenumbernode)
     evaluatordef[:kwargs] = []
 
     # Replace the user-provided function body with the version created by DynamicPPL.
-    evaluatordef[:body] = modelinfo[:body]
+    # NOTE: We need to replace statements of the form `return ...` with
+    # `return DynamicPPL.return_values(..., __varinfo__)` to ensure that the second
+    # element in the returned value is always the most up-to-date `__varinfo__`.
+    # See the docstrings of `replace_returns` and `return_values` for more info.
+    evaluatordef[:body] = replace_returns(make_returns_explicit!(modelinfo[:body]))
 
     ## Build the model function.