Get rid of PrefixContext & make submodels much better

Currently, PrefixContext is used to evaluate submodels. How exactly? -- well, the magic happens in this function.

https://github.com/TuringLang/DynamicPPL.jl/blob/412b2a9f264dca0e4d6634c8b5470050b84bcd1d/src/submodel.jl#L186-L219

TLDR: We evaluate the submodel by wrapping its context in PrefixContext, and evaluating it using the same VarInfo as in the top-level model. PrefixContext hooks into the tilde pipeline where necessary in order to modify what the VarInfo sees. So, for example, a tilde-statement x ~ dist, when passed through PrefixContext, will generate something like tilde_assume!!(..., @varname(prefix.x), ...).

This has so far been fine, but there are two problems with this.

Problems

1. It relies on the submodel being something that can be evaluated with VarInfo. I imagine a day where submodels can be literally anything that obeys a unified API that has (almost) nothing to do with DynamicPPL at all. We could add individual overloads for everything that we might want to use as a submodel, but that's a crappy solution. It'd be much better to have an expected interface that submodel providers can conform to. See below for more info.

2. Template info in the top-level model is lost when entering a submodel. That is to say,

@model function inner()
   a ~ Normal()
end
@model function outer()
   x = zeros(4)
   for i in eachindex(x)
       x[i] ~ to_submodel(inner())
   end
end

Inside the submodel, we don't have access to the shape of x. That's because x is in the top-level model!

In fact, currently on breaking this model is broken. It errors with a very similar error to #1216.

That is easily fixable. It is fixed now by #1222. The problem is that even when we fix #1216, the model will still run, but it won't know the correct size for x, and so it will generate a GrowableArray in the VNT. Everything that is wrong with GrowableArrays will then hit us. For example, using a mixture of linear and Cartesian indices will fail:

@model function outer()
   x = zeros(2, 2)
   x[1] ~ to_submodel(inner())
   x[2,2] ~ to_submodel(inner())
end

(If you only use linear indices then the VNT thinks that it's a vector -- it will run, but results may be subtly different.)

Solutions

Now, how do we fix this?

One possibility is to somehow pass information about the template, down into the submodel. That could conceivably be done by modifying this function:

https://github.com/TuringLang/DynamicPPL.jl/blob/1482ac0a63995a261cbaf42f92cca034a78a3c8f/src/submodel.jl#L174-L182

The fourth, unused, argument is the template, which in the above case, would be the array x.

The difficulty with this is that you have to somehow combine the template info about the array x, with any template info that is inside the submodel itself. That is really tricky. For example, if the inner submodel has index variables itself, then we need to create a "stacked template" that has an array within an array. I think this is impossible to do correctly and generally.

So let's cut the Gordian knot, and stipulate that evaluation inside a submodel is completely independent of evaluation outside a submodel. Prefixing will only happen outside the submodel execution.

From the submodel's point of view, it wouldn't even know that it's being prefixed. Why should it, anyway?! It's downright perplexing that we force the submodel to take on that information which has nothing to do with it.

See the comments in here for more info:

@model function inner()
   a = zeros(1)
   # When executing this tilde statement, we only ever see the VarName a[1].
   # We can pick up the template from `a` as usual. In other words, executing this
   # submodel is *exactly* the same as executing this model on its own.
   a[1] ~ Normal()
end

@model function outer()
   x = zeros(1)

   # Here, we evaluate `inner()` as if it was its own model. We could either use a fresh, empty,
   # VarInfo to do it. Alternatively, if we already have a VarInfo, we could use the value of 
   # `vi[@varname(x[1])]` to do it, which (if the VarInfo was correctly generated), should be a
   # VNT with the correct structure for the submodel.

   x[i] ~ to_submodel(inner())

   # The above will return a tuple of (return value, new_varinfo), where new_varinfo is unprefixed.
   # To regain the original behaviour of prefixed variables, it is trivial to then apply a global prefix
   # to the new VarInfo:
   #
   #   DynamicPPL.templated_setindex!!(vi.values, inner_varinfo.values, @varname(x[i]), x)
   #
   # We also need to merge the accumulators from the inner_varinfo, e.g. summing logp, etc.
   # And then we are all happy and can proceed.
end

Notice that this can be generalised, such that any object called as a submodel only needs to return two things: a return value plus the equivalent of a VarInfo -- i.e., a key-value mapping of unprefixed VarNames. It will then be the top-level model's responsibility to figure out how to use and combine this information with the rest of the model (which is literally what the templated_setindex!! call above does!).

In the case of a DynamicPPL.Model as a submodel, this can be highly optimised as said key-value mapping of VarNames is just a VarInfo obtained from evaluating said model!

There are other details to be worked out around conditioning etc., but I don't think there is any fundamental blocker to this. For example, once Condition and FixedContext use VNTs as their sole data structure, it would be trivial to extract condition_context.values[@varname(x[i])] (which should be a VNT) and condition the inner model on that. Again, the point is the inner model should never need to know what it has been prefixed with.

Start by reading the PrefixContext evaluation code in src/submodel.jl at the linked ranges, along with the discussion of #1222 and templated_setindex!!. Trace how submodels currently use VarInfo, then define the independent submodel return interface and integration behavior described here, including accumulator merging and prefixing; done means PrefixContext is no longer required for submodel execution.