Handle nested PrefixContext

[penelopeysm]

This PR introduces code to better handle PrefixContext nested within arbitrary chains of contexts.

PrefixContext > AbstractContext > PrefixContext

Correctly handles the case where there are multiple PrefixContexts in a chain that are not contiguous, i.e. point (3) in #786.

julia> using DynamicPPL

julia> p3 = PrefixContext{:a}(SamplingContext(PrefixContext{:b}(DefaultContext())))
PrefixContext{:a, SamplingContext{SampleFromPrior, PrefixContext{:b, DefaultContext}, Random.TaskLocalRNG}}(SamplingContext{SampleFromPrior, PrefixContext{:b, DefaultContext}, Random.TaskLocalRNG}(Random.TaskLocalRNG(), SampleFromPrior(), PrefixContext{:b, DefaultContext}(DefaultContext())))

julia> prefix(p3, @varname(x))
b.a.x  # previously this returned a.x

(Note, I still think this should return a.b.x not b.a.x (cf point (2) in #786) but I've refrained from changing this behaviour in this PR because it would be a breaking change.)

AbstractContext > PrefixContext

Correctly handles the case where PrefixContext is not the outermost context, i.e. point (4) in #786.

julia> p4 = SamplingContext(PrefixContext{:a}(DefaultContext()))
SamplingContext{SampleFromPrior, PrefixContext{:a, DefaultContext}, Random.TaskLocalRNG}(Random.TaskLocalRNG(), SampleFromPrior(), PrefixContext{:a, DefaultContext}(DefaultContext()))

julia> prefix(p4, @varname(x))
a.x  # previously this errored

DebugContext > PrefixContext

This also makes sure to prefix variables accordingly inside check_model_and_trace.

Fixes #785:

using DynamicPPL, Distributions

@model function submodel()
  x ~ Normal(0, 1)
  return x
end

@model function parentmodel()
  x1 ~ to_submodel(DynamicPPL.prefix(submodel(), :a), false)
  x2 ~ to_submodel(DynamicPPL.prefix(submodel(), :b), false)
  return x1 + x2
end

DynamicPPL.DebugUtils.check_model_and_trace(parentmodel(); error_on_failure=true)
# used to error, now passes successfully

ValuesAsInModelContext > PrefixContext

Finally, values_as_in_model correctly picks up prefixes now:

Fixes #788

using DynamicPPL, Distributions

@model function submodel()
  x ~ Normal(0, 1)
  return x
end

@model function parentmodel()
  x1 ~ to_submodel(prefix(submodel(), :a), false)
  x2 ~ to_submodel(prefix(submodel(), :b), false)
  return x1 + x2
end

m = parentmodel()
vi = VarInfo(m)
values_as_in_model(m, false, vi)
# OrderedDict{Any, Any} with 2 entries:
#   a.x => 0.640188
#   b.x => 0.477075
# (used to only have a single key `x`)

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 86.17%. Comparing base (727da63) to head (47048c2).
Report is 3 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #787      +/-   ##
==========================================
+ Coverage   86.16%   86.17%   +0.01%     
==========================================
  Files          36       36              
  Lines        4301     4305       +4     
==========================================
+ Hits         3706     3710       +4     
  Misses        595      595              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[coveralls]

Pull Request Test Coverage Report for Build 12970005900

Details

• 6 of 16 (37.5%) changed or added relevant lines in 3 files are covered.
• 4 unchanged lines in 2 files lost coverage.
• Overall coverage increased (+0.01%) to 86.259%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
src/debug_utils.jl	0	10	0.0%

Files with Coverage Reduction	New Missed Lines	%
src/contexts.jl	1	31.77%
src/varinfo.jl	3	84.48%

Totals
Change from base Build 12908348905:	0.01%
Covered Lines:	3710
Relevant Lines:	4301

---

💛 - Coveralls

[penelopeysm]

performance

using DynamicPPL, Distributions, Chairmarks

@model function submodel()
    x ~ Normal(0, 1)
end

@model function parentmodel()
    x1 ~ to_submodel(submodel(), true)
end

m = parentmodel()

@be VarInfo(m)

# Current master
Benchmark: 3827 samples with 2 evaluations
 min    10.667 μs (171 allocs: 8.609 KiB)
 median 11.042 μs (171 allocs: 8.609 KiB)
 mean   12.263 μs (171 allocs: 8.609 KiB, 0.03% gc time)
 max    4.467 ms (171 allocs: 8.609 KiB, 98.57% gc time)

# This PR
julia> @be VarInfo(m)
Benchmark: 3749 samples with 2 evaluations
 min    11.209 μs (171 allocs: 8.609 KiB)
 median 11.729 μs (171 allocs: 8.609 KiB)
 mean   12.650 μs (171 allocs: 8.609 KiB, 0.03% gc time)
 max    3.150 ms (171 allocs: 8.609 KiB, 98.04% gc time)

[torfjelde]

Lovely; thanks @penelopeysm :)

[penelopeysm]

@sunxd3 Will wait for you if you'd like to review, feel free to pass if you don't have time / don't want to:)

[sunxd3]

thanks @penelopeysm, I haven't followed closely with this part of the code, but it looks good to me