From 6aebd9581c55111257ef05ef09edeb5eec111862 Mon Sep 17 00:00:00 2001
From: Penelope Yong <penelopeysm@gmail.com>
Date: Tue, 11 Nov 2025 14:35:15 +0000
Subject: [PATCH 1/8] Remove NodeTrait

---
 src/contexts.jl                | 78 ++++++++++++--------------------
 src/contexts/conditionfix.jl   | 82 ++++++++++------------------------
 src/contexts/default.jl        |  1 -
 src/contexts/init.jl           |  1 -
 src/contexts/prefix.jl         | 17 +++----
 src/contexts/transformation.jl |  1 -
 src/test_utils/contexts.jl     | 19 +++-----
 test/contexts.jl               | 32 +++++--------
 8 files changed, 73 insertions(+), 158 deletions(-)

diff --git a/src/contexts.jl b/src/contexts.jl
index 32a236e8e..5e7465b65 100644
--- a/src/contexts.jl
+++ b/src/contexts.jl
@@ -1,41 +1,19 @@
 """
-    NodeTrait(context)
-    NodeTrait(f, context)
+    AbstractParentContext
 
-Specifies the role of `context` in the context-tree.
-
-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
+An abstract context that has a child context.
 """
-    IsParent
+abstract type AbstractParentContext <: AbstractContext end
 
-Specifies that the context is a parent in the context-tree.
 """
-struct IsParent <: NodeTrait 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.
@@ -60,12 +38,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 +57,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 +79,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
-)
-    return setchildcontext(left, setleafcontext(childcontext(left), right))
-end
-function setleafcontext(::IsParent, ::IsLeaf, left::AbstractContext, right::AbstractContext)
+function setleafcontext(left::AbstractParentContext, right::AbstractParentContext)
     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 +102,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,7 +140,7 @@ 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},
@@ -179,3 +148,12 @@ function tilde_observe!!(
 )
     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
diff --git a/src/contexts/conditionfix.jl b/src/contexts/conditionfix.jl
index d3802de85..333a8940c 100644
--- a/src/contexts/conditionfix.jl
+++ b/src/contexts/conditionfix.jl
@@ -11,7 +11,7 @@ when there are varnames that cannot be represented as symbols, e.g.
 """
 struct ConditionContext{
     Values<:Union{NamedTuple,AbstractDict{<:VarName}},Ctx<:AbstractContext
-} <: AbstractContext
+} <: AbstractParentContext
     values::Values
     context::Ctx
 end
@@ -41,7 +41,6 @@ function Base.show(io::IO, context::ConditionContext)
     return print(io, "ConditionContext($(context.values), $(childcontext(context)))")
 end
 
-NodeTrait(::ConditionContext) = IsParent()
 childcontext(context::ConditionContext) = context.context
 setchildcontext(parent::ConditionContext, child) = ConditionContext(parent.values, child)
 
@@ -76,11 +75,8 @@ Return `true` if `vn` is found in `context` or any of its descendants.
 This is contrast to [`hasconditioned(::AbstractContext, ::VarName)`](@ref) which only checks
 for `vn` in `context`, not recursively checking if `vn` is in any of its descendants.
 """
-function hasconditioned_nested(context::AbstractContext, vn)
-    return hasconditioned_nested(NodeTrait(hasconditioned_nested, context), context, vn)
-end
-hasconditioned_nested(::IsLeaf, context, vn) = hasconditioned(context, vn)
-function hasconditioned_nested(::IsParent, context, vn)
+hasconditioned_nested(context::AbstractContext, vn) = hasconditioned(context, vn)
+function hasconditioned_nested(context::AbstractParentContext, vn)
     return hasconditioned(context, vn) || hasconditioned_nested(childcontext(context), vn)
 end
 function hasconditioned_nested(context::PrefixContext, vn)
@@ -96,15 +92,12 @@ This is contrast to [`getconditioned`](@ref) which only returns the value `vn` i
 not recursively looking into its descendants.
 """
 function getconditioned_nested(context::AbstractContext, vn)
-    return getconditioned_nested(NodeTrait(getconditioned_nested, context), context, vn)
-end
-function getconditioned_nested(::IsLeaf, context, vn)
     return error("context $(context) does not contain value for $vn")
 end
 function getconditioned_nested(context::PrefixContext, vn)
     return getconditioned_nested(collapse_prefix_stack(context), vn)
 end
-function getconditioned_nested(::IsParent, context, vn)
+function getconditioned_nested(context::AbstractParentContext, vn)
     return if hasconditioned(context, vn)
         getconditioned(context, vn)
     else
@@ -113,7 +106,7 @@ function getconditioned_nested(::IsParent, context, vn)
 end
 
 """
-    decondition(context::AbstractContext, syms...)
+    decondition_context(context::AbstractContext, syms...)
 
 Return `context` but with `syms` no longer conditioned on.
 
@@ -121,13 +114,10 @@ Note that this recursively traverses contexts, deconditioning all along the way.
 
 See also: [`condition`](@ref)
 """
-decondition_context(::IsLeaf, context, args...) = context
-function decondition_context(::IsParent, context, args...)
+decondition_context(context::AbstractContext, args...) = context
+function decondition_context(context::AbstractParentContext, args...)
     return setchildcontext(context, decondition_context(childcontext(context), args...))
 end
-function decondition_context(context, args...)
-    return decondition_context(NodeTrait(context), context, args...)
-end
 function decondition_context(context::ConditionContext)
     return decondition_context(childcontext(context))
 end
@@ -160,11 +150,8 @@ Return `NamedTuple` of values that are conditioned on under context`.
 Note that this will recursively traverse the context stack and return
 a merged version of the condition values.
 """
-function conditioned(context::AbstractContext)
-    return conditioned(NodeTrait(conditioned, context), context)
-end
-conditioned(::IsLeaf, context) = NamedTuple()
-conditioned(::IsParent, context) = conditioned(childcontext(context))
+conditioned(::AbstractContext) = NamedTuple()
+conditioned(context::AbstractParentContext) = conditioned(childcontext(context))
 function conditioned(context::ConditionContext)
     # Note the order of arguments to `merge`. The behavior of the rest of DPPL
     # is that the outermost `context` takes precendence, hence when resolving
@@ -176,7 +163,7 @@ function conditioned(context::PrefixContext)
     return conditioned(collapse_prefix_stack(context))
 end
 
-struct FixedContext{Values,Ctx<:AbstractContext} <: AbstractContext
+struct FixedContext{Values,Ctx<:AbstractContext} <: AbstractParentContext
     values::Values
     context::Ctx
 end
@@ -197,7 +184,6 @@ function Base.show(io::IO, context::FixedContext)
     return print(io, "FixedContext($(context.values), $(childcontext(context)))")
 end
 
-NodeTrait(::FixedContext) = IsParent()
 childcontext(context::FixedContext) = context.context
 setchildcontext(parent::FixedContext, child) = FixedContext(parent.values, child)
 
@@ -206,7 +192,7 @@ setchildcontext(parent::FixedContext, child) = FixedContext(parent.values, child
 
 Return `true` if a fixed value for `vn` is found in `context`.
 """
-hasfixed(context::AbstractContext, vn::VarName) = false
+hasfixed(::AbstractContext, ::VarName) = false
 hasfixed(context::FixedContext, vn::VarName) = hasvalue(context.values, vn)
 function hasfixed(context::FixedContext, vns::AbstractArray{<:VarName})
     return all(Base.Fix1(hasvalue, context.values), vns)
@@ -230,11 +216,8 @@ Return `true` if a fixed value for `vn` is found in `context` or any of its desc
 This is contrast to [`hasfixed(::AbstractContext, ::VarName)`](@ref) which only checks
 for `vn` in `context`, not recursively checking if `vn` is in any of its descendants.
 """
-function hasfixed_nested(context::AbstractContext, vn)
-    return hasfixed_nested(NodeTrait(hasfixed_nested, context), context, vn)
-end
-hasfixed_nested(::IsLeaf, context, vn) = hasfixed(context, vn)
-function hasfixed_nested(::IsParent, context, vn)
+hasfixed_nested(context::AbstractContext, vn) = hasfixed(context, vn)
+function hasfixed_nested(context::AbstractParentContext, vn)
     return hasfixed(context, vn) || hasfixed_nested(childcontext(context), vn)
 end
 function hasfixed_nested(context::PrefixContext, vn)
@@ -250,15 +233,12 @@ This is contrast to [`getfixed`](@ref) which only returns the value `vn` in `con
 not recursively looking into its descendants.
 """
 function getfixed_nested(context::AbstractContext, vn)
-    return getfixed_nested(NodeTrait(getfixed_nested, context), context, vn)
-end
-function getfixed_nested(::IsLeaf, context, vn)
     return error("context $(context) does not contain value for $vn")
 end
 function getfixed_nested(context::PrefixContext, vn)
     return getfixed_nested(collapse_prefix_stack(context), vn)
 end
-function getfixed_nested(::IsParent, context, vn)
+function getfixed_nested(context::AbstractParentContext, vn)
     return if hasfixed(context, vn)
         getfixed(context, vn)
     else
@@ -283,7 +263,7 @@ end
 function fix(values::NTuple{<:Any,<:Pair{<:VarName}})
     return fix(DefaultContext(), values)
 end
-fix(context::AbstractContext, values::NamedTuple{()}) = context
+fix(context::AbstractContext, ::NamedTuple{()}) = context
 function fix(context::AbstractContext, values::Union{AbstractDict,NamedTuple})
     return FixedContext(values, context)
 end
@@ -306,13 +286,10 @@ Note that this recursively traverses contexts, unfixing all along the way.
 
 See also: [`fix`](@ref)
 """
-unfix(::IsLeaf, context, args...) = context
-function unfix(::IsParent, context, args...)
+unfix(context::AbstractContext, args...) = context
+function unfix(context::AbstractParentContext, args...)
     return setchildcontext(context, unfix(childcontext(context), args...))
 end
-function unfix(context, args...)
-    return unfix(NodeTrait(context), context, args...)
-end
 function unfix(context::FixedContext)
     return unfix(childcontext(context))
 end
@@ -341,9 +318,8 @@ Return the values that are fixed under `context`.
 Note that this will recursively traverse the context stack and return
 a merged version of the fix values.
 """
-fixed(context::AbstractContext) = fixed(NodeTrait(fixed, context), context)
-fixed(::IsLeaf, context) = NamedTuple()
-fixed(::IsParent, context) = fixed(childcontext(context))
+fixed(::AbstractContext) = NamedTuple()
+fixed(context::AbstractParentContext) = fixed(childcontext(context))
 function fixed(context::FixedContext)
     # Note the order of arguments to `merge`. The behavior of the rest of DPPL
     # is that the outermost `context` takes precendence, hence when resolving
@@ -403,11 +379,8 @@ function collapse_prefix_stack(context::PrefixContext)
     # depth of the context stack.
     return prefix_cond_and_fixed_variables(collapsed, context.vn_prefix)
 end
-function collapse_prefix_stack(context::AbstractContext)
-    return collapse_prefix_stack(NodeTrait(collapse_prefix_stack, context), context)
-end
-collapse_prefix_stack(::IsLeaf, context) = context
-function collapse_prefix_stack(::IsParent, context)
+collapse_prefix_stack(context::AbstractContext) = context
+function collapse_prefix_stack(context::AbstractParentContext)
     new_child_context = collapse_prefix_stack(childcontext(context))
     return setchildcontext(context, new_child_context)
 end
@@ -448,19 +421,10 @@ function prefix_cond_and_fixed_variables(ctx::FixedContext, prefix::VarName)
     prefixed_child_ctx = prefix_cond_and_fixed_variables(childcontext(ctx), prefix)
     return FixedContext(prefixed_vn_dict, prefixed_child_ctx)
 end
-function prefix_cond_and_fixed_variables(c::AbstractContext, prefix::VarName)
-    return prefix_cond_and_fixed_variables(
-        NodeTrait(prefix_cond_and_fixed_variables, c), c, prefix
-    )
-end
-function prefix_cond_and_fixed_variables(
-    ::IsLeaf, context::AbstractContext, prefix::VarName
-)
+function prefix_cond_and_fixed_variables(context::AbstractContext, ::VarName)
     return context
 end
-function prefix_cond_and_fixed_variables(
-    ::IsParent, context::AbstractContext, prefix::VarName
-)
+function prefix_cond_and_fixed_variables(context::AbstractParentContext, prefix::VarName)
     return setchildcontext(
         context, prefix_cond_and_fixed_variables(childcontext(context), prefix)
     )
diff --git a/src/contexts/default.jl b/src/contexts/default.jl
index ec21e1a56..3cafe39f1 100644
--- a/src/contexts/default.jl
+++ b/src/contexts/default.jl
@@ -17,7 +17,6 @@ with `DefaultContext` means 'calculating the log-probability associated with the
 in the `AbstractVarInfo`'.
 """
 struct DefaultContext <: AbstractContext end
-NodeTrait(::DefaultContext) = IsLeaf()
 
 """
     DynamicPPL.tilde_assume!!(
diff --git a/src/contexts/init.jl b/src/contexts/init.jl
index 44dbc5508..83507353f 100644
--- a/src/contexts/init.jl
+++ b/src/contexts/init.jl
@@ -150,7 +150,6 @@ struct InitContext{R<:Random.AbstractRNG,S<:AbstractInitStrategy} <: AbstractCon
         return InitContext(Random.default_rng(), strategy)
     end
 end
-NodeTrait(::InitContext) = IsLeaf()
 
 function tilde_assume!!(
     ctx::InitContext, dist::Distribution, vn::VarName, vi::AbstractVarInfo
diff --git a/src/contexts/prefix.jl b/src/contexts/prefix.jl
index 24615e683..45307874a 100644
--- a/src/contexts/prefix.jl
+++ b/src/contexts/prefix.jl
@@ -13,7 +13,7 @@ unique.
 
 See also: [`to_submodel`](@ref)
 """
-struct PrefixContext{Tvn<:VarName,C<:AbstractContext} <: AbstractContext
+struct PrefixContext{Tvn<:VarName,C<:AbstractContext} <: AbstractParentContext
     vn_prefix::Tvn
     context::C
 end
@@ -23,7 +23,6 @@ function PrefixContext(::Val{sym}, context::AbstractContext) where {sym}
 end
 PrefixContext(::Val{sym}) where {sym} = PrefixContext(VarName{sym}())
 
-NodeTrait(::PrefixContext) = IsParent()
 childcontext(context::PrefixContext) = context.context
 function setchildcontext(ctx::PrefixContext, child::AbstractContext)
     return PrefixContext(ctx.vn_prefix, child)
@@ -37,11 +36,8 @@ Apply the prefixes in the context `ctx` to the variable name `vn`.
 function prefix(ctx::PrefixContext, vn::VarName)
     return AbstractPPL.prefix(prefix(childcontext(ctx), vn), ctx.vn_prefix)
 end
-function prefix(ctx::AbstractContext, vn::VarName)
-    return prefix(NodeTrait(ctx), ctx, vn)
-end
-prefix(::IsLeaf, ::AbstractContext, vn::VarName) = vn
-function prefix(::IsParent, ctx::AbstractContext, vn::VarName)
+prefix(::AbstractContext, vn::VarName) = vn
+function prefix(ctx::AbstractParentContext, vn::VarName)
     return prefix(childcontext(ctx), vn)
 end
 
@@ -72,11 +68,8 @@ function prefix_and_strip_contexts(ctx::PrefixContext, vn::VarName)
     )
     return AbstractPPL.prefix(vn_prefixed, ctx.vn_prefix), child_context_without_prefixes
 end
-function prefix_and_strip_contexts(ctx::AbstractContext, vn::VarName)
-    return prefix_and_strip_contexts(NodeTrait(ctx), ctx, vn)
-end
-prefix_and_strip_contexts(::IsLeaf, ctx::AbstractContext, vn::VarName) = (vn, ctx)
-function prefix_and_strip_contexts(::IsParent, ctx::AbstractContext, vn::VarName)
+prefix_and_strip_contexts(ctx::AbstractContext, vn::VarName) = (vn, ctx)
+function prefix_and_strip_contexts(ctx::AbstractParentContext, vn::VarName)
     vn, new_ctx = prefix_and_strip_contexts(childcontext(ctx), vn)
     return vn, setchildcontext(ctx, new_ctx)
 end
diff --git a/src/contexts/transformation.jl b/src/contexts/transformation.jl
index 5153f7857..c2eee2863 100644
--- a/src/contexts/transformation.jl
+++ b/src/contexts/transformation.jl
@@ -10,7 +10,6 @@ Note that some `AbstractVarInfo` types, must notably `VarInfo`, override the
 how to do the transformation, used by e.g. `SimpleVarInfo`.
 """
 struct DynamicTransformationContext{isinverse} <: AbstractContext end
-NodeTrait(::DynamicTransformationContext) = IsLeaf()
 
 function tilde_assume!!(
     ::DynamicTransformationContext{isinverse},
diff --git a/src/test_utils/contexts.jl b/src/test_utils/contexts.jl
index aae2e4ec6..c48d2ddfd 100644
--- a/src/test_utils/contexts.jl
+++ b/src/test_utils/contexts.jl
@@ -4,11 +4,10 @@
 # Utilities for testing contexts.
 
 # Dummy context to test nested behaviors.
-struct TestParentContext{C<:DynamicPPL.AbstractContext} <: DynamicPPL.AbstractContext
+struct TestParentContext{C<:DynamicPPL.AbstractContext} <: DynamicPPL.AbstractParentContext
     context::C
 end
 TestParentContext() = TestParentContext(DefaultContext())
-DynamicPPL.NodeTrait(::TestParentContext) = DynamicPPL.IsParent()
 DynamicPPL.childcontext(context::TestParentContext) = context.context
 DynamicPPL.setchildcontext(::TestParentContext, child) = TestParentContext(child)
 function Base.show(io::IO, c::TestParentContext)
@@ -25,19 +24,13 @@ This method ensures that `context`
 - Correctly implements the tilde-pipeline.
 """
 function test_context(context::DynamicPPL.AbstractContext, model::DynamicPPL.Model)
-    node_trait = DynamicPPL.NodeTrait(context)
-    if node_trait isa DynamicPPL.IsLeaf
-        test_leaf_context(context, model)
-    elseif node_trait isa DynamicPPL.IsParent
-        test_parent_context(context, model)
-    else
-        error("Invalid NodeTrait: $node_trait")
-    end
+    return test_leaf_context(context, model)
+end
+function test_context(context::DynamicPPL.AbstractParentContext, model::DynamicPPL.Model)
+    return test_parent_context(context, model)
 end
 
 function test_leaf_context(context::DynamicPPL.AbstractContext, model::DynamicPPL.Model)
-    @test DynamicPPL.NodeTrait(context) isa DynamicPPL.IsLeaf
-
     # Note that for a leaf context we can't assume that it will work with an
     # empty VarInfo. (For example, DefaultContext will error with empty
     # varinfos.) Thus we only test evaluation with VarInfos that are already
@@ -57,8 +50,6 @@ function test_leaf_context(context::DynamicPPL.AbstractContext, model::DynamicPP
 end
 
 function test_parent_context(context::DynamicPPL.AbstractContext, model::DynamicPPL.Model)
-    @test DynamicPPL.NodeTrait(context) isa DynamicPPL.IsParent
-
     @testset "get/set leaf and child contexts" begin
         # Ensure we're using a different leaf context than the current.
         leafcontext_new = if DynamicPPL.leafcontext(context) isa DefaultContext
diff --git a/test/contexts.jl b/test/contexts.jl
index 972d833a5..d23b1a6b8 100644
--- a/test/contexts.jl
+++ b/test/contexts.jl
@@ -6,9 +6,6 @@ using DynamicPPL:
     childcontext,
     setchildcontext,
     AbstractContext,
-    NodeTrait,
-    IsLeaf,
-    IsParent,
     contextual_isassumption,
     FixedContext,
     ConditionContext,
@@ -25,22 +22,18 @@ using LinearAlgebra: I
 using Random: Xoshiro
 
 # TODO: Should we maybe put this in DPPL itself?
-function Base.iterate(context::AbstractContext)
-    if NodeTrait(context) isa IsLeaf
-        return nothing
-    end
-
-    return context, context
+function Base.iterate(context::DynamicPPL.AbstractParentContext)
+    return context, childcontext(context)
 end
-function Base.iterate(_::AbstractContext, context::AbstractContext)
-    return _iterate(NodeTrait(context), context)
+function Base.iterate(::DynamicPPL.AbstractContext, state::DynamicPPL.AbstractParentContext)
+    return state, childcontext(state)
 end
-_iterate(::IsLeaf, context) = nothing
-function _iterate(::IsParent, context)
-    child = childcontext(context)
-    return child, child
+function Base.iterate(::DynamicPPL.AbstractContext, state::DynamicPPL.AbstractContext)
+    return state, nothing
+end
+function Base.iterate(::DynamicPPL.AbstractContext, state::Nothing)
+    return nothing
 end
-
 Base.IteratorSize(::Type{<:AbstractContext}) = Base.SizeUnknown()
 Base.IteratorEltype(::Type{<:AbstractContext}) = Base.EltypeUnknown()
 
@@ -347,11 +340,10 @@ Base.IteratorEltype(::Type{<:AbstractContext}) = Base.EltypeUnknown()
         @testset "collapse_prefix_stack" begin
             # Utility function to make sure that there are no PrefixContexts in
             # the context stack.
-            function has_no_prefixcontexts(ctx::AbstractContext)
-                return !(ctx isa PrefixContext) && (
-                    NodeTrait(ctx) isa IsLeaf || has_no_prefixcontexts(childcontext(ctx))
-                )
+            function has_no_prefixcontexts(ctx::AbstractParentContext)
+                return !(ctx isa PrefixContext) && has_no_prefixcontexts(childcontext(ctx))
             end
+            has_no_prefixcontexts(::AbstractContext) = true
 
             # Prefix -> Condition
             c1 = PrefixContext(@varname(a), ConditionContext((c=1, d=2)))

From deac5bc13b61389c365004a483850a59c5d3c7d3 Mon Sep 17 00:00:00 2001
From: Penelope Yong <penelopeysm@gmail.com>
Date: Tue, 11 Nov 2025 14:37:46 +0000
Subject: [PATCH 2/8] Changelog

---
 HISTORY.md | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/HISTORY.md b/HISTORY.md
index 613957c33..840a81d44 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -2,6 +2,20 @@
 
 ## 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.
+
+#### Miscellaneous
+
 Removed the method `returned(::Model, values, keys)`; please use `returned(::Model, ::AbstractDict{<:VarName})` instead.
 
 ## 0.38.9

From 5150cb563d5edc0f8546acdcb0f8bebf33e47f8f Mon Sep 17 00:00:00 2001
From: Penelope Yong <penelopeysm@gmail.com>
Date: Tue, 11 Nov 2025 14:40:39 +0000
Subject: [PATCH 3/8] Fix exports

---
 HISTORY.md        |  3 +++
 src/DynamicPPL.jl | 13 +++++++++----
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/HISTORY.md b/HISTORY.md
index 840a81d44..b21914cf3 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -14,6 +14,9 @@ There should generally be few reasons to define your own parent contexts (the on
 
 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 `condition(model, vals)` and `prefix(model, vals)` instead.
+
 #### Miscellaneous
 
 Removed the method `returned(::Model, values, keys)`; please use `returned(::Model, ::AbstractDict{<:VarName})` instead.
diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
index e66f3fe11..c43bd89d5 100644
--- a/src/DynamicPPL.jl
+++ b/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,

From a392ab74db0aa1257b647ef8cfc30b53ce56f82e Mon Sep 17 00:00:00 2001
From: Penelope Yong <penelopeysm@gmail.com>
Date: Tue, 11 Nov 2025 14:46:10 +0000
Subject: [PATCH 4/8] docs

---
 docs/src/api.md | 37 +++++++++++++++++++++++++++----------
 src/contexts.jl |  6 ++++++
 2 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index bbe39fb73..8d4c73135 100644
--- a/docs/src/api.md
+++ b/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,13 +456,37 @@ 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.
+
+To implement a leaf context, you just need to subtype `AbstractPPL.AbstractContext` and implement the `tilde_assume!!` and `tilde_observe!!` methods for your context.
 
 ```@docs
+AbstractPPL.AbstractContext
 DefaultContext
-PrefixContext
-ConditionContext
 InitContext
+tilde_assume!!
+tilde_observe!!
+```
+
+DynamicPPL has more leaf contexts which are used for internal purposes, but these are the exported ones.
+
+**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
 ```
 
 ### VarInfo initialisation
diff --git a/src/contexts.jl b/src/contexts.jl
index 5e7465b65..12f6b6123 100644
--- a/src/contexts.jl
+++ b/src/contexts.jl
@@ -2,6 +2,12 @@
     AbstractParentContext
 
 An abstract context that has a child context.
+
+Subtypes of `AbstractParentContext` must implement the following interface:
+
+- `DynamicPPL.childcontext(context::AbstractParentContext)`: Return the child context.
+- `DynamicPPL.setchildcontext(parent::AbstractParentContext, child::AbstractContext)`: Reconstruct
+  `parent` but now using `child` as its child context.
 """
 abstract type AbstractParentContext <: AbstractContext end
 

From b13154b5f5e9820245643c802c6b811c1679f0be Mon Sep 17 00:00:00 2001
From: Penelope Yong <penelopeysm@gmail.com>
Date: Tue, 11 Nov 2025 15:03:32 +0000
Subject: [PATCH 5/8] fix a bug

---
 docs/src/api.md  | 19 ++++++++++++++-----
 src/contexts.jl  |  2 +-
 test/contexts.jl | 13 ++++++++-----
 3 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index 8d4c73135..63dafdfca 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -463,19 +463,20 @@ 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.
-
-To implement a leaf context, you just need to subtype `AbstractPPL.AbstractContext` and implement the `tilde_assume!!` and `tilde_observe!!` methods for your context.
+DynamicPPL has more leaf contexts which are used for internal purposes, but these are the two that are exported.
 
 ```@docs
-AbstractPPL.AbstractContext
 DefaultContext
 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!!
 ```
 
-DynamicPPL has more leaf contexts which are used for internal purposes, but these are the exported ones.
-
 **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.
 
@@ -489,6 +490,14 @@ 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.
diff --git a/src/contexts.jl b/src/contexts.jl
index 12f6b6123..c77e7062d 100644
--- a/src/contexts.jl
+++ b/src/contexts.jl
@@ -85,7 +85,7 @@ julia> # Append another parent context.
 ParentContext(ParentContext(ParentContext(DefaultContext())))
 ```
 """
-function setleafcontext(left::AbstractParentContext, right::AbstractParentContext)
+function setleafcontext(left::AbstractParentContext, right::AbstractContext)
     return setchildcontext(left, setleafcontext(childcontext(left), right))
 end
 setleafcontext(::AbstractContext, right::AbstractContext) = right
diff --git a/test/contexts.jl b/test/contexts.jl
index d23b1a6b8..4c7b4c52c 100644
--- a/test/contexts.jl
+++ b/test/contexts.jl
@@ -6,7 +6,9 @@ using DynamicPPL:
     childcontext,
     setchildcontext,
     AbstractContext,
+    AbstractParentContext,
     contextual_isassumption,
+    PrefixContext,
     FixedContext,
     ConditionContext,
     decondition_context,
@@ -22,16 +24,16 @@ using LinearAlgebra: I
 using Random: Xoshiro
 
 # TODO: Should we maybe put this in DPPL itself?
-function Base.iterate(context::DynamicPPL.AbstractParentContext)
+function Base.iterate(context::AbstractParentContext)
     return context, childcontext(context)
 end
-function Base.iterate(::DynamicPPL.AbstractContext, state::DynamicPPL.AbstractParentContext)
+function Base.iterate(::AbstractContext, state::AbstractParentContext)
     return state, childcontext(state)
 end
-function Base.iterate(::DynamicPPL.AbstractContext, state::DynamicPPL.AbstractContext)
+function Base.iterate(::AbstractContext, state::AbstractContext)
     return state, nothing
 end
-function Base.iterate(::DynamicPPL.AbstractContext, state::Nothing)
+function Base.iterate(::AbstractContext, state::Nothing)
     return nothing
 end
 Base.IteratorSize(::Type{<:AbstractContext}) = Base.SizeUnknown()
@@ -340,8 +342,9 @@ Base.IteratorEltype(::Type{<:AbstractContext}) = Base.EltypeUnknown()
         @testset "collapse_prefix_stack" begin
             # Utility function to make sure that there are no PrefixContexts in
             # the context stack.
+            has_no_prefixcontexts(::PrefixContext) = false
             function has_no_prefixcontexts(ctx::AbstractParentContext)
-                return !(ctx isa PrefixContext) && has_no_prefixcontexts(childcontext(ctx))
+                return has_no_prefixcontexts(childcontext(ctx))
             end
             has_no_prefixcontexts(::AbstractContext) = true
 

From 5680e49cc5561a762b4ea41777c0df49b966270f Mon Sep 17 00:00:00 2001
From: Penelope Yong <penelopeysm@gmail.com>
Date: Tue, 11 Nov 2025 15:24:22 +0000
Subject: [PATCH 6/8] Fix doctests

---
 src/contexts.jl              |  2 +-
 src/contexts/conditionfix.jl | 10 +++++++---
 src/model.jl                 |  4 ++--
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/src/contexts.jl b/src/contexts.jl
index c77e7062d..46c5b8855 100644
--- a/src/contexts.jl
+++ b/src/contexts.jl
@@ -26,7 +26,7 @@ effectively updating the child context.
 
 # Examples
 ```jldoctest
-julia> using DynamicPPL: DynamicTransformationContext
+julia> using DynamicPPL: DynamicTransformationContext, ConditionContext
 
 julia> ctx = ConditionContext((; a = 1));
 
diff --git a/src/contexts/conditionfix.jl b/src/contexts/conditionfix.jl
index 333a8940c..7a34db5cb 100644
--- a/src/contexts/conditionfix.jl
+++ b/src/contexts/conditionfix.jl
@@ -42,7 +42,9 @@ function Base.show(io::IO, context::ConditionContext)
 end
 
 childcontext(context::ConditionContext) = context.context
-setchildcontext(parent::ConditionContext, child) = ConditionContext(parent.values, child)
+function setchildcontext(parent::ConditionContext, child::AbstractContext)
+    return ConditionContext(parent.values, child)
+end
 
 """
     hasconditioned(context::AbstractContext, vn::VarName)
@@ -185,7 +187,9 @@ function Base.show(io::IO, context::FixedContext)
 end
 
 childcontext(context::FixedContext) = context.context
-setchildcontext(parent::FixedContext, child) = FixedContext(parent.values, child)
+function setchildcontext(parent::FixedContext, child::AbstractContext)
+    return FixedContext(parent.values, child)
+end
 
 """
     hasfixed(context::AbstractContext, vn::VarName)
@@ -350,7 +354,7 @@ topic](https://turinglang.org/DynamicPPL.jl/previews/PR892/internals/submodel_co
 which explains this in much more detail.
 
 ```jldoctest
-julia> using DynamicPPL: collapse_prefix_stack
+julia> using DynamicPPL: collapse_prefix_stack, PrefixContext, ConditionContext
 
 julia> c1 = PrefixContext(@varname(a), ConditionContext((x=1, )));
 
diff --git a/src/model.jl b/src/model.jl
index ec98b90cd..94fcd9fd4 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -427,7 +427,7 @@ Return the conditioned values in `model`.
 ```jldoctest
 julia> using Distributions
 
-julia> using DynamicPPL: conditioned, contextualize
+julia> using DynamicPPL: conditioned, contextualize, PrefixContext, ConditionContext
 
 julia> @model function demo()
            m ~ Normal()
@@ -770,7 +770,7 @@ Return the fixed values in `model`.
 ```jldoctest
 julia> using Distributions
 
-julia> using DynamicPPL: fixed, contextualize
+julia> using DynamicPPL: fixed, contextualize, PrefixContext
 
 julia> @model function demo()
            m ~ Normal()

From 53e0ef7c11cd128d707fad3f55260487ec55f849 Mon Sep 17 00:00:00 2001
From: Penelope Yong <penelopeysm@gmail.com>
Date: Tue, 11 Nov 2025 15:25:34 +0000
Subject: [PATCH 7/8] Fix test

---
 test/contexts.jl | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/test/contexts.jl b/test/contexts.jl
index 4c7b4c52c..ae7332a43 100644
--- a/test/contexts.jl
+++ b/test/contexts.jl
@@ -27,6 +27,9 @@ using Random: Xoshiro
 function Base.iterate(context::AbstractParentContext)
     return context, childcontext(context)
 end
+function Base.iterate(context::AbstractContext)
+    return context, nothing
+end
 function Base.iterate(::AbstractContext, state::AbstractParentContext)
     return state, childcontext(state)
 end

From e54d3c3e9c015fc3ba15d8bafa2f5af6c3a16890 Mon Sep 17 00:00:00 2001
From: Penelope Yong <penelopeysm@gmail.com>
Date: Tue, 11 Nov 2025 15:27:18 +0000
Subject: [PATCH 8/8] tweak changelog

---
 HISTORY.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/HISTORY.md b/HISTORY.md
index b21914cf3..f181897f7 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -15,7 +15,7 @@ There should generally be few reasons to define your own parent contexts (the on
 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 `condition(model, vals)` and `prefix(model, vals)` instead.
+You should not need to use these directly, please use `AbstractPPL.condition` and `DynamicPPL.prefix` instead.
 
 #### Miscellaneous