Remove NodeTrait (#1133)

* Remove NodeTrait

* Changelog

* Fix exports

* docs

* fix a bug

* Fix doctests

* Fix test

* tweak changelog

Author: penelopeysm

HISTORY.md

@@ -2,6 +2,23 @@
 
 ## 0.39.0
 
+### Breaking changes
+
+#### Parent and leaf contexts
+
+The `DynamicPPL.NodeTrait` function has been removed.
+Instead of implementing this, parent contexts should subtype `DynamicPPL.AbstractParentContext`.
+This is an abstract type which requires you to overload two functions, `DynamicPPL.childcontext` and `DynamicPPL.setchildcontext`.
+
+There should generally be few reasons to define your own parent contexts (the only one we are aware of, outside of DynamicPPL itself, is `Turing.Inference.GibbsContext`), so this change should not really affect users.
+
+Leaf contexts require no changes, apart from a removal of the `NodeTrait` function.
+
+`ConditionContext` and `PrefixContext` are no longer exported.
+You should not need to use these directly, please use `AbstractPPL.condition` and `DynamicPPL.prefix` instead.
+
+#### Miscellaneous
+
 Removed the method `returned(::Model, values, keys)`; please use `returned(::Model, ::AbstractDict{<:VarName})` instead.
 
 ## 0.38.9

docs/src/api.md

@@ -352,13 +352,6 @@ Base.empty!
 SimpleVarInfo
 ```
 
-### Tilde-pipeline
-
-```@docs
-tilde_assume!!
-tilde_observe!!
-```
-
 ### Accumulators
 
 The subtypes of [`AbstractVarInfo`](@ref) store the cumulative log prior and log likelihood, and sometimes other variables that change during executing, in what are called accumulators.
@@ -463,15 +456,48 @@ By default, it does not perform any actual sampling: it only evaluates the model
 If you wish to sample new values, see the section on [VarInfo initialisation](#VarInfo-initialisation) just below this.
 
 The behaviour of a model execution can be changed with evaluation contexts, which are a field of the model.
-Contexts are subtypes of `AbstractPPL.AbstractContext`.
+
+All contexts are subtypes of `AbstractPPL.AbstractContext`.
+
+Contexts are split into two kinds:
+
+**Leaf contexts**: These are the most important contexts as they ultimately decide how model evaluation proceeds.
+For example, `DefaultContext` evaluates the model using values stored inside a VarInfo's metadata, whereas `InitContext` obtains new values either by sampling or from a known set of parameters.
+DynamicPPL has more leaf contexts which are used for internal purposes, but these are the two that are exported.
 
 ```@docs
 DefaultContext
-PrefixContext
-ConditionContext
 InitContext
 ```
 
+To implement a leaf context, you need to subtype `AbstractPPL.AbstractContext` and implement the `tilde_assume!!` and `tilde_observe!!` methods for your context.
+
+```@docs
+tilde_assume!!
+tilde_observe!!
+```
+
+**Parent contexts**: These essentially act as 'modifiers' for leaf contexts.
+For example, `PrefixContext` adds a prefix to all variable names during evaluation, while `ConditionContext` marks certain variables as observed.
+
+To implement a parent context, you have to subtype `DynamicPPL.AbstractParentContext`, and implement the `childcontext` and `setchildcontext` methods.
+If needed, you can also implement `tilde_assume!!` and `tilde_observe!!` for your context.
+This is optional; the default implementation is to simply delegate to the child context.
+
+```@docs
+AbstractParentContext
+childcontext
+setchildcontext
+```
+
+Since contexts form a tree structure, these functions are automatically defined for manipulating context stacks.
+They are mainly useful for modifying the fundamental behaviour (i.e. the leaf context), without affecting any of the modifiers (i.e. parent contexts).
+
+```@docs
+leafcontext
+setleafcontext
+```
+
 ### VarInfo initialisation
 
 The function `init!!` is used to initialise, or overwrite, values in a VarInfo.

src/DynamicPPL.jl

@@ -94,16 +94,21 @@ export AbstractVarInfo,
     values_as_in_model,
     # LogDensityFunction
     LogDensityFunction,
-    # Contexts
+    # Leaf contexts
+    AbstractContext,
     contextualize,
     DefaultContext,
-    PrefixContext,
-    ConditionContext,
+    InitContext,
+    # Parent contexts
+    AbstractParentContext,
+    childcontext,
+    setchildcontext,
+    leafcontext,
+    setleafcontext,
     # Tilde pipeline
     tilde_assume!!,
     tilde_observe!!,
     # Initialisation
-    InitContext,
     AbstractInitStrategy,
     InitFromPrior,
     InitFromUniform,

src/contexts.jl

@@ -1,48 +1,32 @@
 """
-    NodeTrait(context)
-    NodeTrait(f, context)
+    AbstractParentContext
 
-Specifies the role of `context` in the context-tree.
+An abstract context that has a child context.
 
-The officially supported traits are:
-- `IsLeaf`: `context` does not have any decendants.
-- `IsParent`: `context` has a child context to which we often defer.
-  Expects the following methods to be implemented:
-  - [`childcontext`](@ref)
-  - [`setchildcontext`](@ref)
-"""
-abstract type NodeTrait end
-NodeTrait(_, context) = NodeTrait(context)
-
-"""
-    IsLeaf
-
-Specifies that the context is a leaf in the context-tree.
-"""
-struct IsLeaf <: NodeTrait end
-"""
-    IsParent
+Subtypes of `AbstractParentContext` must implement the following interface:
 
-Specifies that the context is a parent in the context-tree.
+- `DynamicPPL.childcontext(context::AbstractParentContext)`: Return the child context.
+- `DynamicPPL.setchildcontext(parent::AbstractParentContext, child::AbstractContext)`: Reconstruct
+  `parent` but now using `child` as its child context.
 """
-struct IsParent <: NodeTrait end
+abstract type AbstractParentContext <: AbstractContext end
 
 """
-    childcontext(context)
+    childcontext(context::AbstractParentContext)
 
 Return the descendant context of `context`.
 """
 childcontext
 
 """
-    setchildcontext(parent::AbstractContext, child::AbstractContext)
+    setchildcontext(parent::AbstractParentContext, child::AbstractContext)
 
 Reconstruct `parent` but now using `child` is its [`childcontext`](@ref),
 effectively updating the child context.
 
 # Examples
 ```jldoctest
-julia> using DynamicPPL: DynamicTransformationContext
+julia> using DynamicPPL: DynamicTransformationContext, ConditionContext
 
 julia> ctx = ConditionContext((; a = 1));
 
@@ -60,12 +44,11 @@ setchildcontext
 """
     leafcontext(context::AbstractContext)
 
-Return the leaf of `context`, i.e. the first descendant context that `IsLeaf`.
+Return the leaf of `context`, i.e. the first descendant context that is not an
+`AbstractParentContext`.
 """
-leafcontext(context::AbstractContext) =
-    leafcontext(NodeTrait(leafcontext, context), context)
-leafcontext(::IsLeaf, context::AbstractContext) = context
-leafcontext(::IsParent, context::AbstractContext) = leafcontext(childcontext(context))
+leafcontext(context::AbstractContext) = context
+leafcontext(context::AbstractParentContext) = leafcontext(childcontext(context))
 
 """
     setleafcontext(left::AbstractContext, right::AbstractContext)
@@ -80,12 +63,10 @@ original leaf context of `left`.
 ```jldoctest
 julia> using DynamicPPL: leafcontext, setleafcontext, childcontext, setchildcontext, AbstractContext, DynamicTransformationContext
 
-julia> struct ParentContext{C} <: AbstractContext
+julia> struct ParentContext{C} <: AbstractParentContext
            context::C
        end
 
-julia> DynamicPPL.NodeTrait(::ParentContext) = DynamicPPL.IsParent()
-
 julia> DynamicPPL.childcontext(context::ParentContext) = context.context
 
 julia> DynamicPPL.setchildcontext(::ParentContext, child) = ParentContext(child)
@@ -104,21 +85,10 @@ julia> # Append another parent context.
 ParentContext(ParentContext(ParentContext(DefaultContext())))
 ```
 """
-function setleafcontext(left::AbstractContext, right::AbstractContext)
-    return setleafcontext(
-        NodeTrait(setleafcontext, left), NodeTrait(setleafcontext, right), left, right
-    )
-end
-function setleafcontext(
-    ::IsParent, ::IsParent, left::AbstractContext, right::AbstractContext
-)
+function setleafcontext(left::AbstractParentContext, right::AbstractContext)
     return setchildcontext(left, setleafcontext(childcontext(left), right))
 end
-function setleafcontext(::IsParent, ::IsLeaf, left::AbstractContext, right::AbstractContext)
-    return setchildcontext(left, setleafcontext(childcontext(left), right))
-end
-setleafcontext(::IsLeaf, ::IsParent, left::AbstractContext, right::AbstractContext) = right
-setleafcontext(::IsLeaf, ::IsLeaf, left::AbstractContext, right::AbstractContext) = right
+setleafcontext(::AbstractContext, right::AbstractContext) = right
 
 """
     DynamicPPL.tilde_assume!!(
@@ -138,10 +108,15 @@ This function should return a tuple `(x, vi)`, where `x` is the sampled value (w
 must be in unlinked space!) and `vi` is the updated VarInfo.
 """
 function tilde_assume!!(
-    context::AbstractContext, right::Distribution, vn::VarName, vi::AbstractVarInfo
+    context::AbstractParentContext, right::Distribution, vn::VarName, vi::AbstractVarInfo
 )
     return tilde_assume!!(childcontext(context), right, vn, vi)
 end
+function tilde_assume!!(
+    context::AbstractContext, ::Distribution, ::VarName, ::AbstractVarInfo
+)
+    return error("tilde_assume!! not implemented for context of type $(typeof(context))")
+end
 
 """
     DynamicPPL.tilde_observe!!(
@@ -171,11 +146,20 @@ This function should return a tuple `(left, vi)`, where `left` is the same as th
 `vi` is the updated VarInfo.
 """
 function tilde_observe!!(
-    context::AbstractContext,
+    context::AbstractParentContext,
     right::Distribution,
     left,
     vn::Union{VarName,Nothing},
     vi::AbstractVarInfo,
 )
     return tilde_observe!!(childcontext(context), right, left, vn, vi)
 end
+function tilde_observe!!(
+    context::AbstractContext,
+    ::Distribution,
+    ::Any,
+    ::Union{VarName,Nothing},
+    ::AbstractVarInfo,
+)
+    return error("tilde_observe!! not implemented for context of type $(typeof(context))")
+end