Avoid pirating AbstractPPL.prefix

Author: penelopeysm

src/compiler.jl

@@ -66,7 +66,7 @@ function isassumption(expr::Union{Expr,Symbol}, left_vn=make_varname_expression(
         # 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)
+        $vn = $(DynamicPPL.maybe_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`,
@@ -454,7 +454,7 @@ function generate_tilde(left, right)
     return quote
         $dist = $right
         $left_vn = $(DynamicPPL.resolve_varnames)($(make_varname_expression(left)), $dist)
-        $vn = $(AbstractPPL.prefix)($left_vn, __model__.prefix)
+        $vn = $(DynamicPPL.maybe_prefix)($left_vn, __model__.prefix)
         $isassumption = $(DynamicPPL.isassumption(left, left_vn))
         if $(DynamicPPL.isfixed(left, vn))
             $left = $(DynamicPPL.getfixed_nested)(__model__.context, $vn)

src/prefix.jl

@@ -1,14 +1,17 @@
 """
-    AbstractPPL.prefix(vn::VarName, ::Nothing)
-AbstractPPL.prefix(::Nothing, vn::VarName)
+    maybe_prefix(inner::Union{Nothing,<:VarName}, outer::Union{Nothing,<:VarName})
 
-Return the original `vn` (i.e., prefixed with nothing).
+Prefix `inner` with the prefix `outer`. Both `inner` and `outer` can be either
+`VarName`s or `Nothing`.
 
-These cases only happen in DynamicPPL and are thus handled here (AbstractPPL's
-definition works only if both arguments are `VarName`s).
+Note that this differs from `AbstractPPL.prefix` in that it handles `nothing` values.
+This can happen e.g. when prefixing a model that is not already prefixed; or when
+executing submodels without automatic prefixing.
 """
-AbstractPPL.prefix(vn::VarName, ::Nothing) = vn
-AbstractPPL.prefix(::Nothing, vn::VarName) = vn
+maybe_prefix(inner::VarName, outer::VarName) = AbstractPPL.prefix(inner, outer)
+maybe_prefix(vn::VarName, ::Nothing) = vn
+maybe_prefix(::Nothing, vn::VarName) = vn
+maybe_prefix(::Nothing, ::Nothing) = nothing
 
 """
     prefix_cond_and_fixed_variables(context::AbstractContext, prefix::VarName)
@@ -94,11 +97,11 @@ julia> rand(prefix(demo(), Val(:my_prefix)))
 ```
 """
 prefix(model::Model, ::Nothing) = model
-function prefix(model::Model, x::VarName)
+function prefix(model::Model, vn::VarName)
     # Add it to the model prefix field
-    new_prefix = AbstractPPL.prefix(model.prefix, x)
+    new_prefix = maybe_prefix(model.prefix, vn)
     # And also make sure to prefix any conditioned and fixed variables stored in the model
-    new_context = prefix_cond_and_fixed_variables(model.context, x)
+    new_context = prefix_cond_and_fixed_variables(model.context, vn)
     return Model(model.f, model.args, model.defaults, new_context, new_prefix)
 end
 prefix(model::Model, ::Val{sym}) where {sym} = prefix(model, VarName{sym}())

src/submodel.jl

@@ -164,7 +164,7 @@ function _evaluate!!(
     vi::AbstractVarInfo,
     parent_context::AbstractContext,
     parent_prefix::Union{Nothing,<:VarName},
-    left_vn::VarName,
+    vn::VarName,
 ) where {M<:Model,AutoPrefix}
     # First, we construct the context to be used when evaluating the submodel. There
     # are several considerations here:
@@ -173,7 +173,9 @@ function _evaluate!!(
     # automatic prefixing if it was requested. (If the prefix was manually applied, then
     # `prefix()` will have been called by the user, and we don't need to do it again.)
     submodel_prefix = if AutoPrefix
-        AbstractPPL.prefix(left_vn, parent_prefix)
+        # Note that by the time we see it here (in `tilde_assume!!`), `vn`
+        # has already prefixed with `parent_prefix`, so no need to re-prefix it
+        vn
     else
         parent_prefix
     end

test/submodels.jl

@@ -133,7 +133,54 @@ using Test
             end
         end
 
-        @testset "Nested submodels" begin
+        @testset "Nested submodels with auto prefix" begin
+            @model function f()
+                x ~ Normal()
+                return y ~ Normal()
+            end
+            @model function g()
+                return b ~ to_submodel(f())
+            end
+            @model function h()
+                return a ~ to_submodel(g())
+            end
+
+            # No conditioning
+            vi = VarInfo(h())
+            @test Set(keys(vi)) == Set([@varname(a.b.x), @varname(a.b.y)])
+            @test getlogjoint(vi) ==
+                logpdf(Normal(), vi[@varname(a.b.x)]) +
+                  logpdf(Normal(), vi[@varname(a.b.y)])
+
+            # Conditioning/fixing at the top level
+            op_h = op(h(), (@varname(a.b.x) => x_val))
+
+            # Conditioning/fixing at the second level
+            op_g = op(g(), (@varname(b.x) => x_val))
+            @model function h2()
+                return a ~ to_submodel(op_g)
+            end
+
+            # Conditioning/fixing at the very bottom
+            op_f = op(f(), (@varname(x) => x_val))
+            @model function g2()
+                return _unused ~ to_submodel(prefix(op_f, :b), false)
+            end
+            @model function h3()
+                return a ~ to_submodel(g2())
+            end
+
+            models = [("top", op_h), ("middle", h2()), ("bottom", h3())]
+            @testset "$name" for (name, model) in models
+                vi = VarInfo(model)
+                @test Set(keys(vi)) == Set([@varname(a.b.y)])
+                @test getlogjoint(vi) == x_logp + logpdf(Normal(), vi[@varname(a.b.y)])
+            end
+        end
+
+        @testset "Nested submodels with manual prefix" begin
+            # Same tests as above, just that the middle layer has manual prefixing
+            # rather than automatic.
             @model function f()
                 x ~ Normal()
                 return y ~ Normal()