Shuffle context code around and remove dead code (#1050)

Author: penelopeysm

src/DynamicPPL.jl

@@ -174,7 +174,11 @@ abstract type AbstractVarInfo <: AbstractModelTrace end
 include("utils.jl")
 include("chains.jl")
 include("contexts.jl")
+include("contexts/default.jl")
 include("contexts/init.jl")
+include("contexts/transformation.jl")
+include("contexts/prefix.jl")
+include("contexts/conditionfix.jl")  # Must come after contexts/prefix.jl
 include("model.jl")
 include("sampler.jl")
 include("varname.jl")
@@ -187,10 +191,8 @@ include("abstract_varinfo.jl")
 include("threadsafe.jl")
 include("varinfo.jl")
 include("simple_varinfo.jl")
-include("context_implementations.jl")
 include("compiler.jl")
 include("pointwise_logdensities.jl")
-include("transforming.jl")
 include("logdensityfunction.jl")
 include("model_utils.jl")
 include("extract_priors.jl")

src/abstract_varinfo.jl

@@ -827,6 +827,27 @@ end
 function link!!(vi::AbstractVarInfo, vns::VarNameTuple, model::Model)
     return link!!(default_transformation(model, vi), vi, vns, model)
 end
+function link!!(t::DynamicTransformation, vi::AbstractVarInfo, model::Model)
+    # Note that in practice this method is only called for SimpleVarInfo, because VarInfo
+    # has a dedicated implementation
+    ctx = DynamicTransformationContext{false}()
+    model = contextualize(model, setleafcontext(model.context, ctx))
+    vi = last(evaluate!!(model, vi))
+    return settrans!!(vi, t)
+end
+function link!!(
+    t::StaticTransformation{<:Bijectors.Transform}, vi::AbstractVarInfo, ::Model
+)
+    b = inverse(t.bijector)
+    x = vi[:]
+    y, logjac = with_logabsdet_jacobian(b, x)
+    # Set parameters and add the logjac term.
+    vi = unflatten(vi, y)
+    if hasacc(vi, Val(:LogJacobian))
+        vi = acclogjac!!(vi, logjac)
+    end
+    return settrans!!(vi, t)
+end
 
 """
     link([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
@@ -846,6 +867,9 @@ end
 function link(vi::AbstractVarInfo, vns::VarNameTuple, model::Model)
     return link(default_transformation(model, vi), vi, vns, model)
 end
+function link(t::DynamicTransformation, vi::AbstractVarInfo, model::Model)
+    return link!!(t, deepcopy(vi), model)
+end
 
 """
     invlink!!([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
@@ -866,23 +890,14 @@ end
 function invlink!!(vi::AbstractVarInfo, vns::VarNameTuple, model::Model)
     return invlink!!(default_transformation(model, vi), vi, vns, model)
 end
-
-# Vector-based ones.
-function link!!(
-    t::StaticTransformation{<:Bijectors.Transform}, vi::AbstractVarInfo, ::Model
-)
-    b = inverse(t.bijector)
-    x = vi[:]
-    y, logjac = with_logabsdet_jacobian(b, x)
-
-    # Set parameters and add the logjac term.
-    vi = unflatten(vi, y)
-    if hasacc(vi, Val(:LogJacobian))
-        vi = acclogjac!!(vi, logjac)
-    end
-    return settrans!!(vi, t)
+function invlink!!(::DynamicTransformation, vi::AbstractVarInfo, model::Model)
+    # Note that in practice this method is only called for SimpleVarInfo, because VarInfo
+    # has a dedicated implementation
+    ctx = DynamicTransformationContext{true}()
+    model = contextualize(model, setleafcontext(model.context, ctx))
+    vi = last(evaluate!!(model, vi))
+    return settrans!!(vi, NoTransformation())
 end
-
 function invlink!!(
     t::StaticTransformation{<:Bijectors.Transform}, vi::AbstractVarInfo, ::Model
 )
@@ -919,6 +934,9 @@ end
 function invlink(vi::AbstractVarInfo, vns::VarNameTuple, model::Model)
     return invlink(default_transformation(model, vi), vi, vns, model)
 end
+function invlink(t::DynamicTransformation, vi::AbstractVarInfo, model::Model)
+    return invlink!!(t, deepcopy(vi), model)
+end
 
 """
     maybe_invlink_before_eval!!([t::Transformation,] vi, model)