Move PrefixContext to a model field

Author: penelopeysm

docs/src/api.md

@@ -476,7 +476,6 @@ Contexts are subtypes of `AbstractPPL.AbstractContext`.
 
 ```@docs
 DefaultContext
-PrefixContext
 ConditionContext
 InitContext
 ```

src/DynamicPPL.jl

@@ -109,7 +109,6 @@ export AbstractVarInfo,
     # Contexts
     contextualize,
     DefaultContext,
-    PrefixContext,
     ConditionContext,
     # Tilde pipeline
     tilde_assume!!,
@@ -184,6 +183,7 @@ include("chains.jl")
 include("contexts.jl")
 include("contexts/init.jl")
 include("model.jl")
+include("prefix.jl")
 include("sampler.jl")
 include("varname.jl")
 include("distribution_wrappers.jl")

src/compiler.jl

@@ -60,11 +60,14 @@ evaluates to a `VarName`, and this will be used in the subsequent checks.
 If `vn` is not specified, `AbstractPPL.varname(expr, need_concretize(expr))` will be
 used in its place.
 """
-function isassumption(expr::Union{Expr,Symbol}, vn=make_varname_expression(expr))
+function isassumption(expr::Union{Expr,Symbol}, left_vn=make_varname_expression(expr))
+    @gensym vn
     return quote
-        if $(DynamicPPL.contextual_isassumption)(
-            __model__.context, $(DynamicPPL.prefix)(__model__.context, $vn)
-        )
+        # TODO(penelopeysm): This re-prefixing seems a bit wasteful. I'd really like
+        # the whole `isassumption` thing to be simplified, though, so I'll
+        # leave it till later.
+        $vn = $(AbstractPPL.prefix)($left_vn, __model__.prefix)
+        if $(DynamicPPL.contextual_isassumption)(__model__.context, $vn)
             # Considered an assumption by `__model__.context` which means either:
             # 1. We hit the default implementation, e.g. using `DefaultContext`,
             #    which in turn means that we haven't considered if it's one of
@@ -78,8 +81,8 @@ function isassumption(expr::Union{Expr,Symbol}, vn=make_varname_expression(expr)
             #    TODO: Support by adding context to model, and use `model.args`
             #    as the default conditioning. Then we no longer need to check `inargnames`
             #    since it will all be handled by `contextual_isassumption`.
-            if !($(DynamicPPL.inargnames)($vn, __model__)) ||
-                $(DynamicPPL.inmissings)($vn, __model__)
+            if !($(DynamicPPL.inargnames)($left_vn, __model__)) ||
+                $(DynamicPPL.inmissings)($left_vn, __model__)
                 true
             else
                 $(maybe_view(expr)) === missing
@@ -99,7 +102,7 @@ isassumption(expr) = :(false)
 
 Return `true` if `vn` is considered an assumption by `context`.
 """
-function contextual_isassumption(context::AbstractContext, vn)
+function contextual_isassumption(context::AbstractContext, vn::VarName)
     if hasconditioned_nested(context, vn)
         val = getconditioned_nested(context, vn)
         # TODO: Do we even need the `>: Missing`, i.e. does it even help the compiler?
@@ -115,9 +118,7 @@ end
 
 isfixed(expr, vn) = false
 function isfixed(::Union{Symbol,Expr}, vn)
-    return :($(DynamicPPL.contextual_isfixed)(
-        __model__.context, $(DynamicPPL.prefix)(__model__.context, $vn)
-    ))
+    return :($(DynamicPPL.contextual_isfixed)(__model__.context, $vn))
 end
 
 """
@@ -448,24 +449,23 @@ function generate_tilde(left, right)
 
     # Otherwise it is determined by the model or its value,
     # if the LHS represents an observation
-    @gensym vn isassumption value dist
+    @gensym left_vn vn isassumption value dist
 
     return quote
         $dist = $right
-        $vn = $(DynamicPPL.resolve_varnames)($(make_varname_expression(left)), $dist)
-        $isassumption = $(DynamicPPL.isassumption(left, vn))
+        $left_vn = $(DynamicPPL.resolve_varnames)($(make_varname_expression(left)), $dist)
+        $vn = $(AbstractPPL.prefix)($left_vn, __model__.prefix)
+        $isassumption = $(DynamicPPL.isassumption(left, left_vn))
         if $(DynamicPPL.isfixed(left, vn))
-            $left = $(DynamicPPL.getfixed_nested)(
-                __model__.context, $(DynamicPPL.prefix)(__model__.context, $vn)
-            )
+            $left = $(DynamicPPL.getfixed_nested)(__model__.context, $vn)
         elseif $isassumption
             $(generate_tilde_assume(left, dist, vn))
         else
-            # If `vn` is not in `argnames`, we need to make sure that the variable is defined.
-            if !$(DynamicPPL.inargnames)($vn, __model__)
-                $left = $(DynamicPPL.getconditioned_nested)(
-                    __model__.context, $(DynamicPPL.prefix)(__model__.context, $vn)
-                )
+            # If `left_vn` is not in `argnames`, we need to make sure that the variable is defined.
+            # (Note: we use the unprefixed `left_vn` here rather than `vn` which will have had
+            # prefixes applied!)
+            if !$(DynamicPPL.inargnames)($left_vn, __model__)
+                $left = $(DynamicPPL.getconditioned_nested)(__model__.context, $vn)
             end
 
             $value, __varinfo__ = $(DynamicPPL.tilde_observe!!)(
@@ -495,6 +495,7 @@ function generate_tilde_assume(left, right, vn)
     return quote
         $value, __varinfo__ = $(DynamicPPL.tilde_assume!!)(
             __model__.context,
+            __model__.prefix,
             $(DynamicPPL.unwrap_right_vn)($(DynamicPPL.check_tilde_rhs)($right), $vn)...,
             __varinfo__,
         )

src/context_implementations.jl

@@ -1,56 +1,30 @@
 # assume
-function tilde_assume!!(context::AbstractContext, right::Distribution, vn, vi)
-    return tilde_assume!!(childcontext(context), right, vn, vi)
+function tilde_assume!!(context::AbstractContext, prefix, right::Distribution, vn, vi)
+    return tilde_assume!!(childcontext(context), prefix, right, vn, vi)
 end
-function tilde_assume!!(::DefaultContext, right::Distribution, vn, vi)
+function tilde_assume!!(::DefaultContext, prefix, right::Distribution, vn, vi)
     y = getindex_internal(vi, vn)
     f = from_maybe_linked_internal_transform(vi, vn, right)
     x, inv_logjac = with_logabsdet_jacobian(f, y)
     vi = accumulate_assume!!(vi, x, -inv_logjac, vn, right)
     return x, vi
 end
-function tilde_assume!!(context::PrefixContext, right::Distribution, vn, vi)
-    # Note that we can't use something like this here:
-    #     new_vn = prefix(context, vn)
-    #     return tilde_assume!!(childcontext(context), right, new_vn, vi)
-    # This is because `prefix` applies _all_ prefixes in a given context to a
-    # variable name. Thus, if we had two levels of nested prefixes e.g.
-    # `PrefixContext{:a}(PrefixContext{:b}(DefaultContext()))`, then the
-    # first call would apply the prefix `a.b._`, and the recursive call
-    # would apply the prefix `b._`, resulting in `b.a.b._`.
-    # This is why we need a special function, `prefix_and_strip_contexts`.
-    new_vn, new_context = prefix_and_strip_contexts(context, vn)
-    return tilde_assume!!(new_context, right, new_vn, vi)
-end
 
 """
-    tilde_assume!!(context, right, vn, vi)
+    tilde_assume!!(context, prefix, right, vn, vi)
 
 Handle assumed variables, e.g., `x ~ Normal()` (where `x` does occur in the model inputs),
 accumulate the log probability, and return the sampled value and updated `vi`.
 """
-function tilde_assume!!(context, right::DynamicPPL.Submodel, vn, vi)
-    return _evaluate!!(right, vi, context, vn)
+function tilde_assume!!(context, prefix, right::DynamicPPL.Submodel, vn, vi)
+    return _evaluate!!(right, vi, context, prefix, vn)
 end
 
 # observe
 function tilde_observe!!(context::AbstractContext, right, left, vn, vi)
     return tilde_observe!!(childcontext(context), right, left, vn, vi)
 end
 
-# `PrefixContext`
-function tilde_observe!!(context::PrefixContext, right, left, vn, vi)
-    # In the observe case, unlike assume, `vn` may be `nothing` if the LHS is a literal
-    # value. For the need for prefix_and_strip_contexts rather than just prefix, see the
-    # comment in `tilde_assume!!`.
-    new_vn, new_context = if vn !== nothing
-        prefix_and_strip_contexts(context, vn)
-    else
-        vn, childcontext(context)
-    end
-    return tilde_observe!!(new_context, right, left, new_vn, vi)
-end
-
 """
     tilde_observe!!(context, right, left, vn, vi)