From 69a12f735bfe0cb44c522af2bb258940c2150271 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Thu, 26 Sep 2024 10:09:33 +0100
Subject: [PATCH] Added impl for `varname_and_value_leaves` that can handle
 containers with multiple varnames typically used in the context of DPPL

---
 src/utils.jl | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/src/utils.jl b/src/utils.jl
index 9ddeb6247..f5057f4d6 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -1028,6 +1028,46 @@ function varname_and_value_leaves(vn::VarName, x)
     return Iterators.map(value, Iterators.flatten(varname_and_value_leaves_inner(vn, x)))
 end
 
+"""
+    varname_and_value_leaves(container)
+
+Return an iterator over all varname-value pairs that are represented by `container`.
+
+This is the same as [`varname_and_value_leaves(vn::VarName, x)`](@ref) but over a container
+containing multiple varnames.
+
+See also: [`varname_and_value_leaves(vn::VarName, x)`](@ref).
+
+# Examples
+```jldoctest varname-and-value-leaves-container
+julia> using DynamicPPL: varname_and_value_leaves
+
+julia> # With an `OrderedDict`
+       dict = OrderedDict(@varname(y) => 1, @varname(z) => [[2.0], [3.0]]);
+
+julia> foreach(println, varname_and_value_leaves(dict))
+(y, 1)
+(z[1][1], 2.0)
+(z[2][1], 3.0)
+
+julia> # With a `NamedTuple`
+       nt = (y = 1, z = [[2.0], [3.0]]);
+
+julia> foreach(println, varname_and_value_leaves(nt))
+(y, 1)
+(z[1][1], 2.0)
+(z[2][1], 3.0)
+```
+"""
+function varname_and_value_leaves(container::OrderedDict)
+    return Iterators.flatten(varname_and_value_leaves(k, v) for (k, v) in container)
+end
+function varname_and_value_leaves(container::NamedTuple)
+    return Iterators.flatten(
+        varname_and_value_leaves(VarName{k}(), v) for (k, v) in pairs(container)
+    )
+end
+
 """
     Leaf{T}