From 1d9f27c26d2dce07a0afbdae20cbdffd4c5e56c4 Mon Sep 17 00:00:00 2001 From: Xianda Sun Date: Sat, 11 Jul 2026 20:06:15 +0100 Subject: [PATCH 1/3] Move VarNamedTuple into AbstractPPL --- HISTORY.md | 6 + Project.toml | 5 +- docs/make.jl | 10 +- docs/src/varnamedtuple.md | 34 + ext/AbstractPPLComponentArraysExt.jl | 55 + ext/AbstractPPLDistributionsExt.jl | 60 +- src/AbstractPPL.jl | 36 + src/varnamedtuple.jl | 196 +++ src/varnamedtuple/getset.jl | 460 ++++++ src/varnamedtuple/macro.jl | 171 ++ src/varnamedtuple/map.jl | 527 ++++++ src/varnamedtuple/partial_array.jl | 804 ++++++++++ src/varnamedtuple/show.jl | 118 ++ src/varnamedtuple/skeleton.jl | 94 ++ src/varnamedtuple/vnt.jl | 250 +++ test/Project.toml | 8 +- test/runtests.jl | 4 +- test/varname/hasvalue.jl | 51 +- test/varnamedtuple.jl | 2232 ++++++++++++++++++++++++++ 19 files changed, 5105 insertions(+), 16 deletions(-) create mode 100644 docs/src/varnamedtuple.md create mode 100644 ext/AbstractPPLComponentArraysExt.jl create mode 100644 src/varnamedtuple.jl create mode 100644 src/varnamedtuple/getset.jl create mode 100644 src/varnamedtuple/macro.jl create mode 100644 src/varnamedtuple/map.jl create mode 100644 src/varnamedtuple/partial_array.jl create mode 100644 src/varnamedtuple/show.jl create mode 100644 src/varnamedtuple/skeleton.jl create mode 100644 src/varnamedtuple/vnt.jl create mode 100644 test/varnamedtuple.jl diff --git a/HISTORY.md b/HISTORY.md index 1ee17908..80f312a3 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -4,6 +4,12 @@ Fixed a Mooncake reverse-mode correctness bug ([Mooncake issue #1238](https://github.com/chalk-lab/Mooncake.jl/issues/1238)): when a problem function closed over differentiable data (e.g. a model capturing observed data reached through an in-place solve), a reused prepared cache accumulated that captured data's cotangent and returned an incorrect gradient after the first evaluation. The Mooncake extension now closes the function and its context into a `NoTangent` target, so the captured data carries no cotangent and reuse is correct. +Moved `VarNamedTuple` and its structured-array support from DynamicPPL into AbstractPPL so +probabilistic programming frontends can share a lossless, `VarName`-keyed parameter +container. The exported API includes `VarNamedTuple`, `@vnt`, `subset`, and the existing +mapping, densification, and skeleton helpers. Template-based insertion remains public as +`templated_setindex!!`. ComponentArrays integration is provided by a package extension. + ## 0.15.3 Added the `of` type system: a self-contained, declarative way to specify the shape, element type, and support of model variables. Construct specifications with the exported `of` function or the exported `@of` macro: diff --git a/Project.toml b/Project.toml index d31b454f..f99d7d62 100644 --- a/Project.toml +++ b/Project.toml @@ -19,6 +19,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" [weakdeps] +ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" @@ -27,6 +28,7 @@ Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [extensions] +AbstractPPLComponentArraysExt = ["ComponentArrays"] AbstractPPLDifferentiationInterfaceExt = ["DifferentiationInterface"] AbstractPPLDistributionsExt = ["Distributions", "LinearAlgebra"] AbstractPPLForwardDiffExt = ["ForwardDiff", "DiffResults"] @@ -37,7 +39,8 @@ AbstractPPLTestExt = ["Test"] ADTypes = "1" AbstractMCMC = "2, 3, 4, 5" Accessors = "0.1" -BangBang = "0.4" +BangBang = "0.4.1" +ComponentArrays = "0.15" DensityInterface = "0.4" DiffResults = "1" DifferentiationInterface = "0.6, 0.7" diff --git a/docs/make.jl b/docs/make.jl index 492c4d09..9f08c3ba 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -10,7 +10,15 @@ DocMeta.setdocmeta!(AbstractPPL, :DocTestSetup, :(using AbstractPPL); recursive= makedocs(; sitename="AbstractPPL", modules=[AbstractPPL, Base.get_extension(AbstractPPL, :AbstractPPLDistributionsExt)], - pages=["index.md", "varname.md", "of.md", "pplapi.md", "evaluators.md", "interface.md"], + pages=[ + "index.md", + "varname.md", + "varnamedtuple.md", + "of.md", + "pplapi.md", + "evaluators.md", + "interface.md", + ], checkdocs=:exports, doctest=false, ) diff --git a/docs/src/varnamedtuple.md b/docs/src/varnamedtuple.md new file mode 100644 index 00000000..1610670a --- /dev/null +++ b/docs/src/varnamedtuple.md @@ -0,0 +1,34 @@ +# VarNamedTuple + +`VarNamedTuple` is a lossless, `VarName`-keyed value container. It preserves nested +property and array structure while allowing probabilistic programming packages to exchange +values without depending on one frontend's model representation. + +Use [`@vnt`](@ref) for literal construction and a `VarName` to access a value. Functional +updates are available through BangBang's `setindex!!` interface. + +```@example varnamedtuple +using AbstractPPL + +values = @vnt begin + location := 1.0 + scale := 2.0 +end +values[@varname(location)] +``` + +## API + +```@docs +VarNamedTuple +AbstractPPL.@vnt +subset +map_pairs!! +map_values!! +apply!! +densify!! +skeleton +NoTemplate +SkipTemplate +AbstractPPL.templated_setindex!! +``` diff --git a/ext/AbstractPPLComponentArraysExt.jl b/ext/AbstractPPLComponentArraysExt.jl new file mode 100644 index 00000000..c857b627 --- /dev/null +++ b/ext/AbstractPPLComponentArraysExt.jl @@ -0,0 +1,55 @@ +module AbstractPPLComponentArraysExt +using AbstractPPL: AbstractPPL +using AbstractPPL.VarNamedTuples: + PartialArray, AllowAll, SetPermissions, _setindex_optic!!, _getindex_optic, make_leaf +using ComponentArrays: ComponentArrays, ComponentVector + +# Helper: convert a Property optic label S to an integer Index optic +function _property_to_index( + template::ComponentVector, optic::AbstractPPL.Property{S} +) where {S} + ax = ComponentArrays.getaxes(template)[1] + idx = first(ax[S].idx) + return AbstractPPL.Index((idx,), NamedTuple(), optic.child) +end + +function AbstractPPL.VarNamedTuples.make_leaf( + value, optic::AbstractPPL.Property{S}, template::ComponentVector +) where {S} + return if optic.child isa AbstractPPL.Iden + index_optic = _property_to_index(template, optic) + make_leaf(value, index_optic, template) + else + # This branch is needed to handle nested axes in ComponentArrays: the idea is that + # if x is e.g. ComponentArray(a=(b=1)) and we are trying to set `x.a.b`, then we + # first index into `x.a` to get the slice of the ComponentArray. The easiest way to + # handle this is to call the default method. + invoke( + make_leaf, + Tuple{Any,AbstractPPL.Property{S},AbstractArray}, + value, + optic, + template, + ) + end +end + +function AbstractPPL.VarNamedTuples._setindex_optic!!( + pa::PartialArray{<:Any,<:Any,<:ComponentVector}, + value, + optic::AbstractPPL.Property{S}, + template, + permissions::SetPermissions=AllowAll(), +) where {S} + index_optic = _property_to_index(pa.data, optic) + return _setindex_optic!!(pa, value, index_optic, template, permissions) +end + +function AbstractPPL.VarNamedTuples._getindex_optic( + pa::PartialArray{<:Any,<:Any,<:ComponentVector}, optic::AbstractPPL.Property{S}, orig_vn +) where {S} + index_optic = _property_to_index(pa.data, optic) + return _getindex_optic(pa, index_optic, orig_vn) +end + +end diff --git a/ext/AbstractPPLDistributionsExt.jl b/ext/AbstractPPLDistributionsExt.jl index 80995365..0216fa78 100644 --- a/ext/AbstractPPLDistributionsExt.jl +++ b/ext/AbstractPPLDistributionsExt.jl @@ -50,6 +50,7 @@ This decision may be revisited in the future. module AbstractPPLDistributionsExt using AbstractPPL: AbstractPPL, VarName, Accessors +using AbstractPPL.VarNamedTuples: VarNamedTuple, PartialArray, _getindex_optic using Distributions: Distributions using LinearAlgebra: Cholesky, LowerTriangular, UpperTriangular @@ -70,7 +71,7 @@ function get_optics(dist::Distributions.LKJCholesky) is_up = dist.uplo == 'U' cartesian_indices = filter(CartesianIndices(size(dist))) do cartesian_index i, j = cartesian_index.I - is_up ? i <= j : i >= j + return is_up ? i <= j : i >= j end # there is an additional layer as we need to access `.L` or `.U` before we # can index into it @@ -302,4 +303,61 @@ function AbstractPPL.getvalue( end end +# A VarNamedTuple already retains its structure, so most distributions can use its direct +# lookup. LKJCholesky is the exception: values may be stored under the `.L` or `.U` optics +# and need the existing dictionary reconstruction path below. +function AbstractPPL.hasvalue( + vals::VarNamedTuple, vn::VarName, ::Distributions.Distribution +) + return haskey(vals, vn) +end +function AbstractPPL.getvalue( + vals::VarNamedTuple, vn::VarName, ::Distributions.Distribution +) + return vals[vn] +end + +function AbstractPPL.hasvalue( + vnt::VarNamedTuple, vn::VarName, dist::Distributions.LKJCholesky +) + haskey(vnt, vn) || return false + val = _getindex_optic(vnt, vn) + if !(val isa VarNamedTuple || val isa PartialArray) + return true + end + + et = val isa VarNamedTuple ? Any : eltype(val) + dval = Dict{VarName,et}() + for key in keys(val) + subvn = if val isa VarNamedTuple + AbstractPPL.prefix(key, vn) + else + AbstractPPL.append_optic(vn, key) + end + dval[subvn] = _getindex_optic(val, key, subvn) + end + return AbstractPPL.hasvalue(dval, vn, dist) +end + +function AbstractPPL.getvalue( + vnt::VarNamedTuple, vn::VarName, dist::Distributions.LKJCholesky +) + val = _getindex_optic(vnt, vn) + if !(val isa VarNamedTuple || val isa PartialArray) + return val + end + + et = val isa VarNamedTuple ? Any : eltype(val) + dval = Dict{VarName,et}() + for key in keys(val) + subvn = if val isa VarNamedTuple + AbstractPPL.prefix(key, vn) + else + AbstractPPL.append_optic(vn, key) + end + dval[subvn] = _getindex_optic(val, key, subvn) + end + return AbstractPPL.getvalue(dval, vn, dist) +end + end diff --git a/src/AbstractPPL.jl b/src/AbstractPPL.jl index f9688fb5..e15ea218 100644 --- a/src/AbstractPPL.jl +++ b/src/AbstractPPL.jl @@ -101,6 +101,42 @@ export AbstractOptic, using Accessors: set export set +""" + subset(value, varnames) + +Return the part of `value` selected by `varnames`. +""" +function subset end + +include("varnamedtuple.jl") +using .VarNamedTuples: + VarNamedTuples, + VarNamedTuple, + map_pairs!!, + map_values!!, + apply!!, + densify!!, + templated_setindex!!, + NoTemplate, + SkipTemplate, + @vnt, + skeleton + +export VarNamedTuple, + @vnt, + map_pairs!!, + map_values!!, + apply!!, + densify!!, + skeleton, + subset, + NoTemplate, + SkipTemplate + +@static if VERSION >= v"1.11.0" + eval(Meta.parse("public VarNamedTuples, templated_setindex!!")) +end + include("of.jl") export of, @of @static if VERSION >= v"1.11.0" diff --git a/src/varnamedtuple.jl b/src/varnamedtuple.jl new file mode 100644 index 00000000..7e46cecc --- /dev/null +++ b/src/varnamedtuple.jl @@ -0,0 +1,196 @@ +module VarNamedTuples + +using BangBang +using ..AbstractPPL +using ..AbstractPPL: AbstractPPL + +export VarNamedTuple, + map_pairs!!, + map_values!!, + apply!!, + templated_setindex!!, + densify!!, + NoTemplate, + SkipTemplate, + @vnt + +""" + NoTemplate + +A singleton struct representing the fact that there is no template for a top-level variable. +When `NoTemplate` is used, several things happen: + + - If you attempt to call `make_leaf` with an Index optic, this will error. + - When recursing into substructures, `NoTemplate` will be propagated. + +Collectively this means that you can only set values for variables that only have Property +optics, unless a template is provided. + +It might seem more idiomatic to use `nothing` or `missing` for this. However, this causes a +bug with BangBang.setindex!!: https://github.com/JuliaFolds2/BangBang.jl/issues/43 so we use +a dedicated struct instead. +""" +struct NoTemplate end + +""" + SkipTemplate{N}(value) + +A struct representing the fact that `value` is the template for the variable `N` levels down +from the top-level variable. In other words, SkipTemplate{0}(value) is equivalent to just +`value`, and SkipTemplate{1}(value) means that `value` is the template for `a` when setting +the variable `@varname(x.a)`. +""" +struct SkipTemplate{N,V} + value::V +end +SkipTemplate{N}(v) where {N} = SkipTemplate{N,typeof(v)}(v) +SkipTemplate{0}(v) = v +SkipTemplate{N}(::NoTemplate) where {N} = NoTemplate() +SkipTemplate{0}(::NoTemplate) = NoTemplate() +# Decrease the skip level: used when recursing inside make_leaf. +function decrease_skip(st::SkipTemplate{N}) where {N} + return SkipTemplate{N - 1}(st.value) +end +# Increase the skip level when a caller prefixes the variable name; the template must be +# wrapped by the corresponding number of SkipTemplates. +SkipTemplate{N}(v::SkipTemplate{M}) where {N,M} = SkipTemplate{N + M}(v.value) +SkipTemplate{0}(v::SkipTemplate{M}) where {M} = SkipTemplate{M}(v.value) + +""" + abstract type SetPermissions end + +An abstract type, passed to `_setindex_optic!!`, that controls whether it is allowed to +create a new value or overwrite an old value in a VarNamedTuple. +""" +abstract type SetPermissions end +""" + AllowAll <: SetPermissions + +Allow the creation of new values, and also allow the overwriting of old values. +""" +struct AllowAll <: SetPermissions end +""" + MustOverwrite(target_vn) <: SetPermissions + +Disallow the creation of a new value for `target_vn`, but allow the overwriting of an old +value. + +This subtype is a performance optimisation: if it is set to `MustOverwrite`, the function +can assume that the key being set already exists in `collection`. This allows skipping some +code paths, which may have a minor benefit at runtime, but more importantly, allows for +better constant propagation and type stability at compile time. + +`permissions` being set to `MustOverwrite` does _not_ guarantee that no new keys will be +added. It only gives the implementation of `_setindex_optic!!` the permission to assume that +the key already exists. Setting it to `MustOverwrite` should be done only when the caller is +sure that the key already exists; anything else is undefined behaviour. +""" +struct MustOverwrite{V<:VarName} <: SetPermissions + target_vn::V +end +struct MustOverwriteError{V<:VarName} <: Exception + target_vn::V + function MustOverwriteError(perm::MustOverwrite) + return new{typeof(perm.target_vn)}(perm.target_vn) + end +end +function Base.showerror(io::IO, e::MustOverwriteError) + # Key doesn't exist yet, so we tried to create it, but we aren't allowed to. + return print( + io, + "MustOverwriteError: Attempted to set a value for $(e.target_vn), but" * + " `permissions=MustOverwrite` was specified. If you did not attempt" * + " to call this function yourself, this likely indicates a bug in" * + " AbstractPPL. Please file an issue at" * + " https://github.com/TuringLang/AbstractPPL.jl/issues.", + ) +end + +""" + MustNotOverwrite <: SetPermissions + +Allow the creation of new values, but disallow the overwriting of old values. +""" +struct MustNotOverwrite{V<:VarName} <: SetPermissions + target_vn::V +end +struct MustNotOverwriteError{V<:VarName} <: Exception + target_vn::V + function MustNotOverwriteError(perm::MustNotOverwrite) + return new{typeof(perm.target_vn)}(perm.target_vn) + end +end +function Base.showerror(io::IO, e::MustNotOverwriteError) + # Key exists already, and we tried to overwrite it, but we aren't allowed to. + return print( + io, + "MustNotOverwriteError: Attempted to set a value for $(e.target_vn), but a value already" * + " existed. This indicates that a value is being set twice (e.g. if" * + " the same variable occurs in a model twice).", + ) +end + +include("varnamedtuple/partial_array.jl") +include("varnamedtuple/vnt.jl") +include("varnamedtuple/getset.jl") +include("varnamedtuple/map.jl") +include("varnamedtuple/show.jl") +include("varnamedtuple/macro.jl") +include("varnamedtuple/skeleton.jl") + +""" + NamedTuple(vnt::VarNamedTuple) + +Convert a `VarNamedTuple` to a standard `NamedTuple`, provided all keys in the +`VarNamedTuple` are `VarName`s with top-level symbols. If any key is a `VarName` +with a non-identity optic (e.g., `@varname(x.a)` or `@varname(x[1])`), this will +throw an `ArgumentError`. + +# Examples + +```jldoctest +julia> using AbstractPPL, BangBang + +julia> vnt = VarNamedTuple(); vnt = setindex!!(vnt, 10, @varname(x)) +VarNamedTuple +└─ x => 10 + +julia> NamedTuple(vnt) +(x = 10,) + +julia> vnt2 = setindex!!(vnt, 20, @varname(y.a)) +VarNamedTuple +├─ x => 10 +└─ y => VarNamedTuple + └─ a => 20 + +julia> NamedTuple(vnt2) +ERROR: ArgumentError: Cannot convert VarNamedTuple containing non-identity VarNames to NamedTuple. To create a NamedTuple, all keys in the VarNamedTuple must be top-level symbols. +[...] +``` +""" +@generated function Base.NamedTuple(vnt::VarNamedTuple{names,vals}) where {names,vals} + if isempty(names) + return :(NamedTuple()) + end + nt = Expr(:tuple) + for (n, v) in zip(names, vals.parameters) + if v <: VarNamedTuple || + VarNamedTuple <: v || + v <: PartialArray || + PartialArray <: v + throw( + ArgumentError( + "Cannot convert VarNamedTuple containing non-identity VarNames to NamedTuple. To create a NamedTuple, all keys in the VarNamedTuple must be top-level symbols.", + ), + ) + end + push!(nt.args, :($n = vnt.data.$n)) + end + return nt +end + +AbstractPPL.hasvalue(vnt::VarNamedTuple, vn::VarName) = haskey(vnt, vn) +AbstractPPL.getvalue(vnt::VarNamedTuple, vn::VarName) = vnt[vn] + +end diff --git a/src/varnamedtuple/getset.jl b/src/varnamedtuple/getset.jl new file mode 100644 index 00000000..356039a9 --- /dev/null +++ b/src/varnamedtuple/getset.jl @@ -0,0 +1,460 @@ +# We define our own getindex, setindex!!, and haskey functions, which we use to +# get/set/check values in VarNamedTuple and PartialArray. We do this because we want to be +# able to override their behaviour for some types exported from elsewhere without type +# piracy. This is needed because +# 1. We would want to index into things with lenses (from AbstractPPL.jl) using getindex and +# setindex!!, but AbstractPPL does not define these methods. +# 2. We would want `haskey` to fall back onto `checkbounds` when called on Base.Arrays. + +const IndexWithoutChild = AbstractPPL.Index{<:Tuple,<:NamedTuple,AbstractPPL.Iden} + +_unimplemented() = error("Not implemented") + +""" + AbstractPPL.VarNamedTuples._getindex_optic( + collection, optic::AbstractPPL.AbstractOptic, orig_vn::VarName + ) + AbstractPPL.VarNamedTuples._getindex_optic(collection, vn::VarName) + +Access the value in `collection` at the location specified by the given `optic`. If a `VarName` +is provided, it is first converted to an optic using `AbstractPPL.varname_to_optic`. + +Here, `collection` can be either a `VarNamedTuple` or a `PartialArray`, or a leaf value stored +within one of these. + +This is semantically similar to `AbstractPPL.getvalue` but is specialised for `VarNamedTuple` +and `PartialArray`, and skips a number of checks that are unnecessary here. + +Note that it is only valid to index into a `VarNamedTuple` with a `Property` optic, and a +`PartialArray` with an `Index` optic. Other combinations are not valid. When we have reached +the leaf of the VNT i.e. a value, we could still handle pure `Index` optics if the value is +an `AbstractArray`, but otherwise the only valid optic is `Iden`. + +`orig_vn` is used to keep track of the original VarName used to index into a VarNamedTuple, +and is only for error reporting purposes. +""" +function _getindex_optic(vnt::VarNamedTuple, vn::VarName) + return _getindex_optic(vnt, AbstractPPL.varname_to_optic(vn), vn) +end +function _getindex_optic(vnt::VarNamedTuple, vn::VarName, orig_vn) + return _getindex_optic(vnt, AbstractPPL.varname_to_optic(vn), orig_vn) +end + +@inline _getindex_optic(@nospecialize(x::Any), ::AbstractPPL.Iden, orig_vn) = x +@inline _getindex_optic(x::Any, o::AbstractPPL.AbstractOptic, orig_vn) = o(x) +function _getindex_optic( + vnt::VarNamedTuple, optic::AbstractPPL.Property{S}, orig_vn +) where {S} + if !haskey(vnt.data, S) + throw(KeyError(orig_vn)) + end + return _getindex_optic(getindex(vnt.data, S), optic.child, orig_vn) +end +function _getindex_optic(pa::PartialArray, optic::AbstractPPL.Index, orig_vn) + coptic = AbstractPPL.concretize_top_level(optic, pa.data) + child_value = + if _is_multiindex(pa, coptic.ix...; coptic.kw...) && + coptic.child isa AbstractPPL.Index + # getindex can error if we try to access an index that is masked. However, it + # may be that the child index subsets into only valid indices (i.e. unmasked + # parts). For example, consider the case where `pa[1]` is valid but `pa[2]` is + # masked. Then `pa[1:2][1]` is valid, even though `pa[1:2]` would error. + _subset_partialarray(pa, coptic.ix...; coptic.kw...) + else + getindex(pa, coptic.ix...; coptic.kw...) + end + return _getindex_optic(child_value, optic.child, orig_vn) +end +function _getindex_optic(arr::AbstractArray, optic::IndexWithoutChild, orig_vn) + coptic = AbstractPPL.concretize_top_level(optic, arr) + return Base.getindex(arr, coptic.ix...; coptic.kw...) +end + +# Fallback for everything else. +@inline _haskey_optic(x::Any, o::AbstractPPL.AbstractOptic) = AbstractPPL.canview(o, x) +function _haskey_optic(vnt::VarNamedTuple, name::VarName) + return _haskey_optic(vnt, AbstractPPL.varname_to_optic(name)) +end +@inline _haskey_optic(@nospecialize(::Any), ::AbstractPPL.Iden) = true +@inline _haskey_optic(::VarNamedTuple, ::AbstractPPL.Index) = false +function _haskey_optic(vnt::VarNamedTuple, optic::AbstractPPL.Property{S}) where {S} + return Base.haskey(vnt.data, S) && _haskey_optic(getindex(vnt.data, S), optic.child) +end +function _haskey_optic(pa::PartialArray, ::AbstractPPL.Iden) + # no masked indices + all(pa.mask) || return false + # no partial ALBs + et = eltype(pa.data) + if et <: ArrayLikeBlock || ArrayLikeBlock <: et + return _can_get_arraylikeblock(pa.data) + else + # no ALBs at all, so it's fine + return true + end +end +function _haskey_optic(pa::PartialArray, optic::AbstractPPL.Index) + # check the top level Index + Base.haskey(pa, optic.ix...; optic.kw...) || return false + # recurse if necessary + return if optic.child isa AbstractPPL.Iden + true + elseif _is_multiindex(pa, optic.ix...; optic.kw...) + _haskey_optic(_subset_partialarray(pa, optic.ix...; optic.kw...), optic.child) + else + _haskey_optic(getindex(pa, optic.ix...; optic.kw...), optic.child) + end +end +function _haskey_optic(arr::AbstractArray, optic::IndexWithoutChild) + # Note that this call to `checkbounds` can error, although it is technically out of our + # hands: it depends on how the provider of the AbstractArray has implemented + # checkbounds. For example, DimArray can error here: + # https://github.com/rafaqz/DimensionalData.jl/issues/1156. But that is not our job to fix + # -- it should be done upstream -- hence we just forward the indices. + return checkbounds(Bool, arr, optic.ix...; optic.kw...) +end + +""" + _setindex_optic!!(collection, value, optic, template, permissions) + +Like `setindex!!`, but special-cased for `VarNamedTuple` and `PartialArray` to recurse into +nested structures. + +The `permissions` argument controls the behaviour depending on whether `optic` already +exists in `vnt` or does not yet exist. See the docstrings of `AllowAll`, `MustOverwrite`, +and `MustNotOverwrite` for details. +""" +@inline function _setindex_optic!!( + @nospecialize(::Any), + value, + ::AbstractPPL.Iden, + @nospecialize(::Any), + perms::SetPermissions, +) + perms isa MustOverwrite && throw(MustOverwriteError(perms)) + return value +end +function _setindex_optic!!( + arr::AbstractArray, value, optic::IndexWithoutChild, template, perms::SetPermissions +) + perms isa MustNotOverwrite && throw(MustNotOverwriteError(perms)) + coptic = AbstractPPL.concretize_top_level(optic, arr) + return BangBang.setindex!!(arr, value, coptic.ix...; coptic.kw...) +end +function _setindex_optic!!( + nt::NamedTuple{names}, + value, + ::AbstractPPL.Property{S,AbstractPPL.Iden}, + template, + perms::SetPermissions, +) where {names,S} + # TODO(penelopeysm): Should this method be generalised to all structs? + if S in names && perms isa MustNotOverwrite + throw(MustNotOverwriteError(perms)) + end + return BangBang.setproperty!!(nt, S, value) +end + +struct SharedGetProperty{S} end +(::SharedGetProperty)(::NoTemplate) = NoTemplate() +(::SharedGetProperty)(t::SkipTemplate{N}) where {N} = decrease_skip(t) +# These two are unlikely to be hit unless someone manually specifies nothing or missing as a +# template, but we can handle them gracefully anyway +(::SharedGetProperty)(::Nothing) = NoTemplate() +(::SharedGetProperty)(::Missing) = NoTemplate() +@generated function (::SharedGetProperty{S})(x::VarNamedTuple{names}) where {S,names} + return (S in names) ? :(x.data.$S) : :(NoTemplate()) +end +@generated function (::SharedGetProperty{S})(x::NamedTuple{names}) where {S,names} + return (S in names) ? :(x.$S) : :(NoTemplate()) +end +# This method handles other structs. We can assume that the property must exist, since +# `x.a ~ dist` would error if `x` did not have a property `a`. +(::SharedGetProperty{S})(x) where {S} = getproperty(x, S) + +function _setindex_optic!!( + pa::PartialArray, value, optic::AbstractPPL.Index, template, permissions::SetPermissions +) + need_merge = false + coptic = AbstractPPL.concretize_top_level(optic, pa.data) + # We might be attempting to set a slice into a PartialArray{GrowableArray} that + # doesn't yet have enough indices for that slice. Expand it if so. + pa = grow_to_indices!!(pa, coptic.ix...; coptic.kw...) + + is_multiindex = _is_multiindex(template, coptic.ix...; coptic.kw...) + + if permissions isa MustNotOverwrite && optic.child isa AbstractPPL.Iden + if any(view(pa.mask, coptic.ix...; coptic.kw...)) + throw(MustNotOverwriteError(permissions)) + end + end + + sub_value = if optic.child isa AbstractPPL.Iden + # Skip recursion + value + else + child_template = if template === NoTemplate() + NoTemplate() + elseif template isa SkipTemplate + decrease_skip(template) + elseif is_multiindex + Base.getindex(template, coptic.ix...; coptic.kw...) + elseif isassigned(template, coptic.ix...; coptic.kw...) + # Single-index, but we should check that there is actual data there before + # calling getindex, as otherwise it will error + Base.getindex(template, coptic.ix...; coptic.kw...) + else + NoTemplate() + end + + # TODO(penelopeysm): This check to haskey() will check *all* the indices, it only + # returns true if they are all filled. This is probably not really correct, and is + # why we need to have the hacky need_merge workaround just below this (because that + # catches the case where *some* of the indices are filled). We should rethink the + # logic in this section. + if Base.haskey(pa, coptic.ix...; coptic.kw...) + # The PartialArray already contains an unmasked value at this index. We need to + # set that value in place there. + _setindex_optic!!( + Base.getindex(pa, coptic.ix...; coptic.kw...), + value, + coptic.child, + child_template, + permissions, + ) + else + # No existing value. + # If we aren't allowed to make one, we should error. + if permissions isa MustOverwrite + throw(MustOverwriteError(permissions)) + end + # Otherwise we can go ahead and make it + if any(view(pa.mask, coptic.ix...; coptic.kw...)) + # NOTE: This is a VERY subtle case, which can happen when you are setting + # multiple indices at once, but some of them were masked. When they are + # masked, it will cause haskey to return false, so we can't go into the + # previous branch. However, if we naively call make_leaf to create the + # sub-value and then setindex it inside the PartialArray, it will overwrite + # ALL the indices with the new leaf value, which will overwrite any + # previously active values! To avoid this, we will set a flag to indicate + # that we can create a new leaf, but at the end of the function we need to + # MERGE the new leaf into the existing PartialArray instead of overwriting + # it. + # This situation can happen e.g. with + # using AbstractPPL + # vnt = VarNamedTuple() + # x = zeros(2) + # vnt = AbstractPPL.templated_setindex!!(vnt, 1.0, @varname(x[1:2][1]), x) + # vnt = AbstractPPL.templated_setindex!!(vnt, 2.0, @varname(x[1:2][2]), x) + need_merge = true + end + # No new data but we are allowed to create it. + make_leaf(value, coptic.child, child_template) + end + end + + # If sub_value is a GrowableArray, we need to make sure it is grown to the right size to + # fit into the indices specified by `coptic`. This is the same logic as in + # `make_leaf_multiindex`. Again, if there is no GrowableArray underpinning sub_value + # then grow_to_indices!! is a no-op so won't hurt. + grown_sub_value = if is_multiindex && sub_value isa PartialArray + grow_to_indices!!(sub_value, coptic.ix...; coptic.kw...) + else + sub_value + end + + # In the merge path, some indices in the slice already have data but not all of them + # (haskey returned false but any(mask) was true). If MustNotOverwrite is set, check + # that the new sub-value doesn't overlap with existing data at the specific sub-indices. + if need_merge && permissions isa MustNotOverwrite && grown_sub_value isa PartialArray + existing_mask = view(pa.mask, coptic.ix...; coptic.kw...) + if any(existing_mask .& grown_sub_value.mask) + throw(MustNotOverwriteError(permissions)) + end + end + + return if need_merge + new_pa = BangBang.setindex!!(copy(pa), grown_sub_value, coptic.ix...; coptic.kw...) + _merge(pa, new_pa, Val(false)) + else + BangBang.setindex!!(pa, grown_sub_value, coptic.ix...; coptic.kw...) + end +end + +function _setindex_optic!!( + vnt::VarNamedTuple{names}, + value, + optic::AbstractPPL.Property{S}, + template, + permissions=AllowAll(), +) where {names,S} + if S in names && optic.child isa AbstractPPL.Iden && permissions isa MustNotOverwrite + throw(MustNotOverwriteError(permissions)) + end + sub_value = if optic.child isa AbstractPPL.Iden + # Skip recursion + value + else + child_template = SharedGetProperty{S}()(template) + if S in names + # Data already exists; we need to recurse into it + _setindex_optic!!(vnt.data[S], value, optic.child, child_template, permissions) + elseif permissions isa MustOverwrite + throw(MustOverwriteError(permissions)) + else + # AllowAll or MustNotOverwrite -- in either case we can just create a new leaf + make_leaf(value, optic.child, child_template) + end + end + return VarNamedTuple(merge(vnt.data, NamedTuple{(S,)}((sub_value,)))) +end + +""" + make_leaf(value, optic, template) + +Make a new leaf node for a VarNamedTuple. + +This is the function that sets any `optic` that is a `Property` to be stored as a +`VarNamedTuple`, any `Index` to be stored as a `PartialArray`, and other `Iden` optics to be +stored as raw values. It is the link that joins `VarNamedTuple` and `PartialArray` together. + +A lot of the complexity in the VarNamedTuple code comes from this function, so its branches +are documented inline. Read those comments before modifying it. +""" +@inline make_leaf(@nospecialize(value::Any), ::AbstractPPL.Iden, ::Any) = value +function make_leaf(value, optic::AbstractPPL.Property{S}, template) where {S} + sub_value = if optic.child isa AbstractPPL.Iden + value + else + child_template = SharedGetProperty{S}()(template) + make_leaf(value, optic.child, child_template) + end + return VarNamedTuple(NamedTuple{(S,)}((sub_value,))) +end + +# This function handles Index optics. Since Index optics can represent either single-element +# indexing or multi-element indexing (e.g., slices, arrays of indices), and their +# implementation can be somewhat different, we dispatch to separate functions for each case. +function make_leaf(value, optic::AbstractPPL.Index, template::PartialArray) + # If the template is a PA, use its data as the template. + return make_leaf(value, optic, template.data) +end +function make_leaf( + value, + optic::AbstractPPL.Index, + template::Union{AbstractArray,NoTemplate,SkipTemplate,Missing}, +) + # First we need to resolve any dynamic indices, since _is_multiindex doesn't work with + # them. This also helpfully catches errors if there is a dynamic index and a suitable + # template is not provided (e.g., if someone tries to set `x[end]` without a template). + coptic = AbstractPPL.concretize_top_level(optic, template) + return if _is_multiindex(template, coptic.ix...; coptic.kw...) + make_leaf_multiindex(value, coptic, template) + else + make_leaf_singleindex(value, coptic, template) + end +end +function make_leaf(value, ::AbstractPPL.Index, template::Any) + # Fallback method; this is hit for other indexable things that we don't know how to deal + # with. + return error( + "VarNamedTuple currently does not support indexed values within any container that is not an `AbstractArray` (got type $(typeof(template))).", + ) +end + +function make_sub_value(value, coptic::AbstractPPL.Index, template) + return if coptic.child isa AbstractPPL.Iden + value + else + child_template = if template isa NoTemplate + NoTemplate() + elseif template isa SkipTemplate + decrease_skip(template) + elseif template isa AbstractArray + # Note template can't be a PartialArray as that would have been unwrapped in + # another make_leaf method. + getindex(template, coptic.ix...; coptic.kw...) + else + NoTemplate() + end + make_leaf(value, coptic.child, child_template) + end +end + +# This is the easy case, and is directly analogous to the Property optic case above. +# The optic is called `coptic` to indicate that it has already been concretised +# by the caller. +function make_leaf_singleindex(value, coptic::AbstractPPL.Index, template) + sub_value = make_sub_value(value, coptic, template) + # `sub_value` will become a single element inside the PartialArray that we create. To + # ensure type stability, we want to make sure that the PartialArray is created with the + # correct eltype to begin with, otherwise setindex!! may become type unstable. + pa_eltype = typeof(sub_value) + pa_data = if template isa AbstractArray + if pa_eltype <: eltype(template) + # Reuse the eltype of template: it might be more abstract than sub_value_type, + # but that's fine since it will mimic the semantics of what the user passed in. + similar(template) + else + # Use the eltype of sub_value. + similar(template, pa_eltype) + end + else + # If no template was provided, or an incorrectly typed template was provided, we + # have to make a GrowableArray. No need to use `kw` since make_leaf already errors + # if there are uninterpretable keyword indices. + template_sz = get_maximum_size_from_indices(coptic.ix...) + _warn_growable_array_creation(template_sz) + GrowableArray(Array{pa_eltype}(undef, template_sz)) + end + pa_mask = similar(pa_data, Bool) + fill!(pa_mask, false) + pa = PartialArray(pa_data, pa_mask) + return BangBang.setindex!!(pa, sub_value, coptic.ix...; coptic.kw...) +end + +# This is more complex. +function make_leaf_multiindex(value, coptic::AbstractPPL.Index, template) + sub_value = make_sub_value(value, coptic, template) + + # Firstly, we need to make sure that sub_value has *exactly* the right size to fit into + # the indices specified by `coptic`. This might not always be the case. Consider + # _[2:3][1] without a template -- the inner `make_leaf` call will create a GrowableArray + # of size 1 (because that's the minimum size it infers from the inner indices), but we + # actually need a slice of length 2. (If there's a template, we don't have this problem, + # because the inner make_leaf will use the template to set the correct size of 2.) Note + # that `grow_to_indices!!` is a no-op for any PartialArray that doesn't contain a + # GrowableArray. + grown_sub_value = if sub_value isa PartialArray + grow_to_indices!!(sub_value, coptic.ix...; coptic.kw...) + else + sub_value + end + + # `sub_value` is not just one element, but is actually a slice inside the PartialArray + # that we create. To add to that, it might not just be a slice, but it might be an + # ArrayLikeBlock that is set at multiple indices. + pa_eltype = if sub_value isa AbstractArray || sub_value isa PartialArray + eltype(sub_value) + else + idx_size_type = Dims{_ndims(template, coptic.ix...;coptic.kw...)} + ArrayLikeBlock{typeof(sub_value),typeof(coptic.ix),typeof(coptic.kw),idx_size_type} + end + + # The rest is the same as the single-index case. + pa_data = if template isa AbstractArray + if pa_eltype <: eltype(template) + similar(template) + else + similar(template, pa_eltype) + end + else + # No template, or incorrectly typed template + template_sz = get_maximum_size_from_indices(coptic.ix...) + _warn_growable_array_creation(template_sz) + GrowableArray(Array{pa_eltype}(undef, template_sz)) + end + pa_mask = similar(pa_data, Bool) + fill!(pa_mask, false) + pa = PartialArray(pa_data, pa_mask) + return BangBang.setindex!!(pa, grown_sub_value, coptic.ix...; coptic.kw...) +end diff --git a/src/varnamedtuple/macro.jl b/src/varnamedtuple/macro.jl new file mode 100644 index 00000000..a0da404e --- /dev/null +++ b/src/varnamedtuple/macro.jl @@ -0,0 +1,171 @@ +""" + @vnt begin ... end + +Construct a `VarNamedTuple` from a block of assignments. Each assignment should be of the form +`var := value`, where `var` is a variable name. This is best illustrated by +example: + +```jldoctest +julia> using AbstractPPL + +julia> @vnt begin + a := 1 + b := 2 + end +VarNamedTuple +├─ a => 1 +└─ b => 2 +``` + +You can set entirely arbitrary variables: + +```jldoctest; setup=:(using AbstractPPL) +julia> @vnt begin + a.b.c.d.e := "hello" + end +VarNamedTuple +└─ a => VarNamedTuple + └─ b => VarNamedTuple + └─ c => VarNamedTuple + └─ d => VarNamedTuple + └─ e => "hello" +``` + +For variables that have indexing, it is often necessary to provide a template, so that the +VNT 'knows' what kind of array is being used to store the values, and can set the values in +the appropriate places (consider e.g. OffsetArrays where `x[1]` may not mean what it usually +does). + +This is done by inserting a `@template` macro call in the block. The `@template` macro +accepts whitespace-separated arguments, which must either be + +- a single symbol (e.g. `@template x`), in which case the template is the value of `x` + (and `x` must already be defined in the current scope); or +- an assignment of the form `y = expr`, in which case the template for `y` is the value of + `expr`. In this case `y` does not need to be defined in the current scope, but any symbols + referenced in `expr` must be. + +For example: + +```jldoctest; setup=:(using AbstractPPL) +julia> x = zeros(5); outside_y = zeros(3, 3); + +julia> @vnt begin + @template x y=outside_y + x[1] := 1.0 + y[1, 1] := 2.0 + end +VarNamedTuple +├─ x => PartialArray size=(5,) data::Vector{Float64} +│ └─ (1,) => 1.0 +└─ y => PartialArray size=(3, 3) data::Matrix{Float64} + └─ (1, 1) => 2.0 +``` + +!!! note + You can use any expression in `@template y=expr`, even a function call that is + completely contained within the macro (e.g. `@template y=zeros(3, 3)`). The macro makes + sure that `expr` is only evaluated once, so there is no performance penalty to doing + this. + +If no template is provided, the VNT will use a `GrowableArray`. This can produce correct +results in simple cases, but is not recommended for general use. Please see the +VarNamedTuple documentation for more details. + +```jldoctest; setup=:(using AbstractPPL) +julia> @vnt begin + # No template provided. + x[1] := 1.0 + y[1, 1] := 2.0 + end +┌ Warning: Creating a growable `Base.Array` of dimension 1 to store values. This may not match the actual type or size of the actual `AbstractArray` represented by the `VarNamedTuple`. +│ +│ If this is not the type or size that you expect, please see: https://turinglang.org/docs/uri/growablearray +└ @ AbstractPPL.VarNamedTuples /path/to/AbstractPPL.jl/src/varnamedtuple/partial_array.jl:823 +┌ Warning: Creating a growable `Base.Array` of dimension 2 to store values. This may not match the actual type or size of the actual `AbstractArray` represented by the `VarNamedTuple`. +│ +│ If this is not the type or size that you expect, please see: https://turinglang.org/docs/uri/growablearray +└ @ AbstractPPL.VarNamedTuples /path/to/AbstractPPL.jl/src/varnamedtuple/partial_array.jl:823 +VarNamedTuple +├─ x => PartialArray size=(1,) data::AbstractPPL.VarNamedTuples.GrowableArray{Float64, 1} +│ └─ (1,) => 1.0 +└─ y => PartialArray size=(1, 1) data::AbstractPPL.VarNamedTuples.GrowableArray{Float64, 2} + └─ (1, 1) => 2.0 +``` +""" +macro vnt(expr) + return _vnt(expr) +end + +# This helper was copied from DynamicPPL's compiler. Keeping it local avoids coupling the +# container implementation to DynamicPPL's model compiler. +_get_top_sym(expr::Symbol) = expr +function _get_top_sym(expr::Expr) + if Meta.isexpr(expr, :ref) + return _get_top_sym(expr.args[1]) + elseif Meta.isexpr(expr, :.) + return _get_top_sym(expr.args[1]) + else + error("unreachable") + end +end + +function _vnt(input) + Meta.isexpr(input, :block) || + error("`@vnt` expects a block expression (e.g. `@vnt begin ... end`)") + @gensym vnt + symbols_to_templates = Dict{Symbol,Union{Expr,Symbol}}() + output = Expr(:block) + push!(output.args, :($vnt = VarNamedTuple())) + for expr in input.args + if expr isa LineNumberNode + push!(output.args, expr) + elseif Meta.isexpr(expr, :macrocall) && expr.args[1] == Symbol("@template") + # Extract all the templated symbols + for arg in expr.args[2:end] + if arg isa LineNumberNode + continue + elseif arg isa Symbol + # e.g. @template x + if arg in keys(symbols_to_templates) + error("duplicate template definition for symbol: $arg") + end + symbols_to_templates[arg] = esc(arg) + elseif Meta.isexpr(arg, :(=)) + # e.g. @template y = x + sym, template_expr = arg.args + if sym in keys(symbols_to_templates) + error("duplicate template definition for symbol $sym") + end + # evaluate the template expression one time so that we don't + # reevaluate it every time we set a value in the VNT + new_sym = gensym() + push!(output.args, :($new_sym = $(esc(template_expr)))) + symbols_to_templates[sym] = new_sym + else + error("unexpected argument to `@template`: $arg") + end + end + elseif Meta.isexpr(expr, :(:=)) + lhs, rhs = expr.args + vn = AbstractPPL.varname(lhs, false) + top_level_sym = _get_top_sym(lhs) + if top_level_sym in keys(symbols_to_templates) + push!( + output.args, + :( + $vnt = AbstractPPL.templated_setindex!!( + $vnt, $(esc(rhs)), $vn, $(symbols_to_templates[top_level_sym]) + ) + ), + ) + else + push!(output.args, :($vnt = BangBang.setindex!!($vnt, $(esc(rhs)), $vn))) + end + else + error("unexpected expression in `@vnt` block: $expr") + end + end + push!(output.args, :($vnt)) + return output +end diff --git a/src/varnamedtuple/map.jl b/src/varnamedtuple/map.jl new file mode 100644 index 00000000..1178c080 --- /dev/null +++ b/src/varnamedtuple/map.jl @@ -0,0 +1,527 @@ +function Base.merge(a::VarNamedTuple, b::VarNamedTuple, cs::VarNamedTuple...) + return merge(merge(a, b), cs...) +end +Base.merge(x1::VarNamedTuple, x2::VarNamedTuple) = _merge(x1, x2, Val(true)) +Base.merge(x1::VarNamedTuple) = x1 + +# This needs to be a generated function for type stability. +@generated function _merge( + vnt1::VarNamedTuple{names1}, vnt2::VarNamedTuple{names2}, ::Val{true} +) where {names1,names2} + all_names = union(names1, names2) + exs = Expr[] + push!(exs, :(data = (;))) + for name in all_names + val_expr = if name in names1 && name in names2 + :(_merge(vnt1.data.$name, vnt2.data.$name, Val(true))) + elseif name in names1 + :(vnt1.data.$name) + else + :(vnt2.data.$name) + end + push!(exs, :(data = merge(data, NamedTuple{($(QuoteNode(name)),)}(($val_expr,))))) + end + push!(exs, :(return VarNamedTuple(data))) + return Expr(:block, exs...) +end + +""" + subset(vnt::VarNamedTuple, vns) + +Create a new `VarNamedTuple` containing only the variables subsumed by ones in `vns`. +""" +function AbstractPPL.subset(parent_vnt::VarNamedTuple, vns) + return mapfoldl( + identity, + function (acc_vnt, pair) + name, value = pair + return if any(vn -> subsumes(vn, name), vns) + templated_setindex!!( + acc_vnt, value, name, parent_vnt.data[AbstractPPL.getsym(name)] + ) + else + acc_vnt + end + end, + parent_vnt; + init=VarNamedTuple(), + ) +end + +""" + apply!!(func, vnt::VarNamedTuple, name::VarName) + +Apply `func` to the subdata at `name` in `vnt`, and set the result back at `name`. + +Like `map_values!!`, but only for a single `VarName`. + +```jldoctest +julia> using AbstractPPL: VarNamedTuple + +julia> using BangBang: setindex!! + +julia> using AbstractPPL.VarNamedTuples: apply!! + +julia> vnt = VarNamedTuple() +VarNamedTuple() + +julia> vnt = setindex!!(vnt, [1, 2, 3], @varname(a)) +VarNamedTuple +└─ a => [1, 2, 3] + +julia> apply!!(x -> x .+ 1, vnt, @varname(a)) +VarNamedTuple +└─ a => [2, 3, 4] +``` +""" +function apply!!(func, vnt::VarNamedTuple, name::VarName) + if !haskey(vnt, name) + throw(KeyError(name)) + end + subdata = _getindex_optic(vnt, name) + new_subdata = func(subdata) + # The MustOverwrite is a performance optimisation: Since we've already checked that + # the key exists, we know that no new fields will be created. + return _setindex_optic!!( + vnt, + new_subdata, + AbstractPPL.varname_to_optic(name), + NoTemplate(), + MustOverwrite(name), + ) +end + +""" + _map_pairs_recursive!!(func, x, vn) + +Call `func` on `vn => x`, except if `x` is a `VarNamedTuple` or `PartialArray`, in which +case call `_map_pairs_recursive!!` recursively on all their elements, updating `vn` with the +right prefix. + +This is the internal implementation of `map_pairs!!`, but because it has a method defined +for literally every type in existence, we hide it behind the interface of the more +discriminating `map_pairs!!`. It makes the implementation a bit simpler, compared to +checking element types within `map_pairs!!` itself. +""" +_map_pairs_recursive!!(pairfunc, x, vn) = pairfunc(vn => x) +_map_values_recursive!!(func, x) = func(x) + +# TODO(mhauru) The below is type unstable for some complex VarNames. My example case +# for which type stability fails is @varname(e.f[3].g.h[2].i). I don't understand this +# well, but I think it's just because constant propagation gives up at some point, and fails +# to go through the lines that figure out `new_et`. I could be wrong. I tried fixing this by +# lifting the first three lines of the function into a generated function, but that seems +# to run into trouble when trying to call Core.Compiler.return_type recursively on the same +# function. An earlier implementation of this function that only operated on the values, +# not on pairs of key => value, was type stable (presumably because it was a bit easier on +# constant propagation). +function _map_pairs_recursive!!(func, pa::PartialArray, vn) + return if eltype(pa) <: ArrayLikeBlock || ArrayLikeBlock <: eltype(pa) + # There are (or might be) some ALBs + _map_pairs_recursive_pa_alb!!(func, pa, vn) + else + # There are definitely no ALBs + _map_pairs_recursive_pa_noalb!!(func, pa, vn) + end +end +function _map_values_recursive!!(func, pa::PartialArray) + return if eltype(pa) <: ArrayLikeBlock || ArrayLikeBlock <: eltype(pa) + # There are (or might be) some ALBs + _map_values_recursive_pa_alb!!(func, pa) + else + # There are definitely no ALBs + _map_values_recursive_pa_noalb!!(func, pa) + end +end + +function _map_pairs_recursive_pa_alb!!(pairfunc, pa::PartialArray, vn) + # Ask the compiler to infer the return type of applying func recursively to eltype(pa). + et = eltype(pa) + index_type = AbstractPPL.Index{NTuple{ndims(pa),Int},@NamedTuple{},AbstractPPL.Iden} + new_vn_type = Core.Compiler.return_type( + AbstractPPL.append_optic, Tuple{typeof(vn),index_type} + ) + new_et = Core.Compiler.return_type( + Tuple{typeof(_map_pairs_recursive!!),typeof(pairfunc),et,new_vn_type} + ) + # Confusingly, sometimes the compiler just gives up and returns Union{}. We'll bail it + # out. TODO(penelopeysm): I have noticed that sometimes return_type is not amazingly + # good at inferring types, especially when there are ArrayLikeBlocks -- Not sure if + # there is a better way to hack into the type system here. + new_et = new_et == Union{} ? Any : new_et + new_data = if new_et <: et + # We can reuse the existing data array. + pa.data + else + # We need to allocate a new data array. + similar(pa.data, new_et) + end + # If there are any ArrayLikeBlocks, we will need to modify the mask during the + # iteration, so make a copy. The alternative, which avoids mutating, is to store a Dict + # of already-visited ArrayLikeBlocks along with the result of applying `func` to them, + # but that is more expensive. + mask = copy(pa.mask) + for i in CartesianIndices(mask) + if mask[i] + val = pa.data[i] + if val isa ArrayLikeBlock + # Create the new value, set it at all the places where it needs to be set, + # and mark those places as visited. + new_vn = AbstractPPL.append_optic(vn, AbstractPPL.Index(val.ix, val.kw)) + new_val = _map_pairs_recursive!!(pairfunc, val, new_vn) + new_data[val.ix..., val.kw...] .= new_val + mask[val.ix..., val.kw...] .= false + else + new_vn = AbstractPPL.append_optic(vn, AbstractPPL.Index(Tuple(i), (;))) + new_data[i] = _map_pairs_recursive!!(pairfunc, val, new_vn) + end + end + end + # The above type inference may be overly conservative, so we concretise the eltype. + return _concretise_eltype!!(PartialArray(new_data, pa.mask)) +end +function _map_values_recursive_pa_alb!!(func, pa::PartialArray) + et = eltype(pa) + new_et = Core.Compiler.return_type( + Tuple{typeof(_map_values_recursive!!),typeof(func),et} + ) + new_et = new_et == Union{} ? Any : new_et + new_data = if new_et <: et + pa.data + else + similar(pa.data, new_et) + end + mask = copy(pa.mask) + for i in CartesianIndices(mask) + if mask[i] + val = pa.data[i] + if val isa ArrayLikeBlock + # Create the new value, set it at all the places where it needs to be set, + # and mark those places as visited. + new_val = _map_values_recursive!!(func, val) + new_data[val.ix..., val.kw...] .= new_val + mask[val.ix..., val.kw...] .= false + else + new_data[i] = _map_values_recursive!!(func, val) + end + end + end + return _concretise_eltype!!(PartialArray(new_data, pa.mask)) +end + +function _map_pairs_recursive_pa_noalb!!(pairfunc, pa::PartialArray, vn) + # Don't have to faff with ALBs. Just map `func` over all elements that are set. + mask = pa.mask + ci = CartesianIndices(mask) + ixs_subset = @view ci[mask] + data_subset = @view pa.data[mask] + index_type = AbstractPPL.Index{NTuple{ndims(pa),Int},@NamedTuple{},AbstractPPL.Iden} + new_vn_type = Core.Compiler.return_type( + AbstractPPL.append_optic, Tuple{typeof(vn),index_type} + ) + et = eltype(pa) + new_et = Core.Compiler.return_type( + Tuple{typeof(_map_pairs_recursive!!),typeof(pairfunc),eltype(pa),new_vn_type} + ) + function inner(data_i, ix_i) + new_vn = AbstractPPL.append_optic(vn, AbstractPPL.Index(Tuple(ix_i), (;))) + return _map_pairs_recursive!!(pairfunc, data_i, new_vn) + end + return if new_et <: et + # Can write into existing data array. + map!(inner, view(pa.data, mask), data_subset, ixs_subset) + pa + else + # Need to allocate a new data array. + new_pa_data = similar(pa.data, new_et) + map!(inner, view(new_pa_data, mask), data_subset, ixs_subset) + _concretise_eltype!!(PartialArray(new_pa_data, pa.mask)) + end +end +function _map_values_recursive_pa_noalb!!(func, pa::PartialArray) + mask = pa.mask + et = eltype(pa) + new_et = Core.Compiler.return_type( + Tuple{typeof(_map_values_recursive!!),typeof(func),eltype(pa)} + ) + return if new_et <: et + # Can write into existing data array. + new_pa_data = view(pa.data, mask) + # According to the docstring of `map!` on 1.11, "unexpected behaviour" may happen if + # the input and output alias each other. But the docstring on 1.12 says that you can + # do this, so... + # https://docs.julialang.org/en/v1/base/collections/#Julia-1.12-9b4e28751024c91d + map!(Base.Fix1(_map_values_recursive!!, func), new_pa_data, new_pa_data) + pa + else + # Need to allocate a new data array. + new_pa_data = similar(pa.data, new_et) + map!( + Base.Fix1(_map_values_recursive!!, func), + view(new_pa_data, mask), + view(pa.data, mask), + ) + _concretise_eltype!!(PartialArray(new_pa_data, pa.mask)) + end +end + +function _map_pairs_recursive!!(pairfunc, alb::ArrayLikeBlock, vn) + new_block = pairfunc(vn => alb.block) + return ArrayLikeBlock(new_block, alb.ix, alb.kw, alb.index_size) +end +function _map_values_recursive!!(func, alb::ArrayLikeBlock) + new_block = _map_values_recursive!!(func, alb.block) + return ArrayLikeBlock(new_block, alb.ix, alb.kw, alb.index_size) +end + +@generated function _map_pairs_recursive!!( + pairfunc, vnt::VarNamedTuple{Names}, vn::T +) where {Names,T} + exs = Expr[] + for name in Names + push!( + exs, + :(_map_pairs_recursive!!( + pairfunc, + vnt.data.$name, + AbstractPPL.append_optic(vn, AbstractPPL.Property{$(QuoteNode(name))}()), + )), + ) + end + return quote + return VarNamedTuple(NamedTuple{Names}(($(exs...),))) + end +end +@generated function _map_values_recursive!!(func, vnt::VarNamedTuple{Names}) where {Names} + exs = Expr[] + for name in Names + push!(exs, :(_map_values_recursive!!(func, vnt.data.$name))) + end + return quote + return VarNamedTuple(NamedTuple{Names}(($(exs...),))) + end +end + +""" + map_pairs!!(func, vnt::VarNamedTuple) + +Apply `func` to all key => value pairs of `vnt`, in place if possible. + +`func` should accept a pair of `VarName` and value, and return the new value to be set. +""" +@generated function map_pairs!!(func, vnt::VarNamedTuple{Names}) where {Names} + exs = Expr[] + for name in Names + push!( + exs, + :(_map_pairs_recursive!!(func, vnt.data.$name, VarName{$(QuoteNode(name))}())), + ) + end + return quote + return VarNamedTuple(NamedTuple{Names}(($(exs...),))) + end +end + +""" + map_values!!(func, vnt::VarNamedTuple) + +Apply `func` to elements of `vnt`, in place if possible. +""" +@generated function map_values!!(func, vnt::VarNamedTuple{Names}) where {Names} + # NOTE(penelopeysm): The definition of this function, and the functions it recurses + # into, is almost entirely the same as for map_pairs!!, except that we don't generate + # and pass around the VarName. This leads to a LOT of code duplication, unfortunately. + # One could get rid of this by defining: + # + # map_values!!(func, vnt) = map_pairs!!((pair) -> func(pair.second), vnt) + # + # Unfortunately, that does cause substantial performance regressions on functions such + # as unflatten!! (differences of 10x have been seen on some models). So do make sure + # you benchmark any changes to this function carefully. + # + # We could also use @eval to reduce the code duplication, but I personally don't think + # it's worth the added complexity right now. I would be open to a PR that changed this, + # though. + exs = Expr[] + for name in Names + push!(exs, :(_map_values_recursive!!(func, vnt.data.$name))) + end + return quote + return VarNamedTuple(NamedTuple{Names}(($(exs...),))) + end +end + +""" + mapreduce(f, op, vnt::VarNamedTuple; init) + +Apply `f` to all elements of `vnt`, and reduce the results using `op`, starting from `init`. + +The order is the same as in `mapfoldl`, i.e. left-associative with `init` as the +left-most value. + +`init` is a keyword argument to conform to the usual `mapreduce` interface in Base, but it +is not optional. + +`f` op` should accept pairs of `varname => value`. +""" +@generated function Base.mapreduce( + f, op, vnt::VarNamedTuple{Names}; init::InitType=nothing +) where {Names,InitType} + # NOTE(penelopeysm): I tried doing a separate implementation for mapreduce and + # mapreduce_pairs, just like for map and map_pairs, but it actually led to slowdowns. + if InitType === Nothing + return quote + throw( + ArgumentError( + "mapreduce without init is not implemented for VarNamedTuple." + ), + ) + end + end + + exs = Expr[:(result = init)] + for name in Names + push!( + exs, + quote + result = _mapreduce_recursive( + f, op, vnt.data.$name, VarName{$(QuoteNode(name))}(), result + ) + end, + ) + end + push!(exs, :(return result)) + return Expr(:block, exs...) +end + +# Our mapreduce is always left-associative. +Base.mapfoldl(f, op, vnt::VarNamedTuple; init=nothing) = mapreduce(f, op, vnt; init=init) + +_mapreduce_recursive(f, op, x, vn, init) = op(init, f(vn => x)) +function _mapreduce_recursive(f, op, alb::ArrayLikeBlock, vn, init) + return op(init, f(vn => alb.block)) +end + +# As above but with a prefix VarName `vn`. +@generated function _mapreduce_recursive( + f, op, vnt::VarNamedTuple{Names}, vn, init +) where {Names} + exs = Expr[:(result = init)] + for name in Names + push!( + exs, + quote + result = _mapreduce_recursive( + f, + op, + vnt.data.$name, + AbstractPPL.prefix(VarName{$(QuoteNode(name))}(), vn), + result, + ) + end, + ) + end + push!(exs, :(return result)) + return Expr(:block, exs...) +end + +function _mapreduce_recursive(f, op, pa::PartialArray, vn, init) + result = init + # If there are any ArrayLikeBlocks, we will need to modify the mask during the + # iteration, so make a copy. The alternative, which avoids mutating, is to store a Set + # of already-visited ArrayLikeBlocks, but that is more expensive. + mask = if eltype(pa) <: ArrayLikeBlock || ArrayLikeBlock <: eltype(pa) + copy(pa.mask) + else + pa.mask + end + for i in CartesianIndices(mask) + if mask[i] + val = pa.data[i] + is_alb = val isa ArrayLikeBlock + if is_alb + # Don't visit the same ArrayLikeBlock multiple times. + mask[val.ix..., val.kw...] .= false + end + optic = if is_alb + AbstractPPL.Index(val.ix, val.kw) + else + AbstractPPL.Index(Tuple(i), (;)) + end + new_vn = AbstractPPL.append_optic(vn, optic) + result = _mapreduce_recursive(f, op, pa.data[i], new_vn, result) + end + end + return result +end + +""" + densify!!(vnt::VarNamedTuple) + +Convert any `PartialArray`s in `vnt` where all elements are filled to just normal arrays. +""" +@generated function densify!!(vnt::VarNamedTuple{Names}) where {Names} + exs = Expr[] + for name in Names + push!(exs, :(densify!!(vnt.data.$name))) + end + return :(VarNamedTuple(NamedTuple{Names}(($(exs...),)))) +end +# The method on PartialArrays is inherently type unstable because it has to check +# `all(pa.mask)`, which cannot be inferred at compile time. +function densify!!(pa::PartialArray) + # Don't densify GrowableArrays, since we can't be sure that that's what the user really + # intended. + if pa.data isa GrowableArray + return pa + end + + # We can't densify a PA whose elements are ALBs or VNTs (densifying would lose + # VNT-based indexing / ALB semantics). But we still recursively densify the elements + # themselves, so that any nested PAs inside VNTs get densified. + et = eltype(pa) + has_albs = (et <: ArrayLikeBlock || ArrayLikeBlock <: et) + has_vnts = (et <: VarNamedTuple || VarNamedTuple <: et) + if has_albs || has_vnts + new_data = map(densify!!, pa.data) + return PartialArray(new_data, pa.mask) + end + + # If the mask is not all true, then it's not dense. The function past this point is no + # longer type stable. + if !all(pa.mask) + return pa + end + + # If we reach here, then we can make a serious attempt at densifying. + has_nested_pas = (et <: PartialArray || PartialArray <: et) + return if has_nested_pas + # Attempt to recursively densify the elements. + new_data = map(densify!!, pa.data) + if !(any(x -> x isa PartialArray, new_data)) + # Inner PAs were densified successfully, so we can densify this PA too + new_data + else + pa + end + else + # No nested PAs, so we can just densify this PA directly. + pa.data + end +end +@inline densify!!(@nospecialize(x::Any)) = x + +# TODO(mhauru) We could try to keep the return types of these more tight, rather than always +# return the same, abstract element type. Would that be better? It would be faster in some +# cases, but would be less consistent, and could result in a lot of allocations in the +# mapreduce, as the element type is gradually expanded. +Base.keys(vnt::VarNamedTuple) = mapreduce(first, push!, vnt; init=VarName[]) +Base.values(vnt::VarNamedTuple) = mapreduce(pair -> pair.second, push!, vnt; init=Any[]) + +function Base.length(vnt::VarNamedTuple) + len = 0 + for subdata in vnt.data + len += subdata isa VarNamedTuple || subdata isa PartialArray ? length(subdata) : 1 + end + return len +end diff --git a/src/varnamedtuple/partial_array.jl b/src/varnamedtuple/partial_array.jl new file mode 100644 index 00000000..298f64f5 --- /dev/null +++ b/src/varnamedtuple/partial_array.jl @@ -0,0 +1,804 @@ +# Some utilities for checking what sort of indices we are dealing with. +# The non-generated function implementations of these would be +# _has_colon(::T) where {T<:Tuple} = any(x <: Colon for x in T.parameters) +# However, constant propagation sometimes fails if the index tuple is too big (e.g. length +# 4), so we play it safe and use generated functions. Constant propagating these is +# important, because many functions choose different paths based on their values, which +# would lead to type instability if they were only evaluated at runtime. +@generated function _has_colon_or_dynamicindex(::T) where {T<:Tuple} + for x in T.parameters + if x <: Colon || x <: AbstractPPL.DynamicIndex + return :(return true) + end + end + return :(return false) +end + +""" + _merge(x1, x2, recursive=Val(true)) + +Merge two values `x1` and `x2`. If `recursive isa Val{true}`, additionally recurses into +elements of `x1` and `x2` to merge those. + +Unlike `Base.merge`, this function is defined for all types, and by default returns the +second argument. It is overridden for `PartialArray` and `VarNamedTuple`, since they are +nested containers, and the `recursive` argument allows us to control whether we want to +merge its nested values as well. +""" +_merge(_, x2, _) = x2 + +""" + ArrayLikeBlock{T,I,N,S} + +A wrapper for non-array blocks stored in `PartialArray`s. + +When setting a value in a `PartialArray` over a range of indices, if the value being set +is not itself an `AbstractArray`, but has a well-defined size, we wrap it in an +`ArrayLikeBlock`, which records both the value and the indices it was set with. + +When getting values from a `PartialArray`, if any of the requested indices correspond to +an `ArrayLikeBlock`, we check that the requested indices match the ones used to set the +value. If they do, we return the underlying block, otherwise we throw an error. +""" +struct ArrayLikeBlock{T,I,N,S} + block::T + ix::I + kw::N + index_size::S +end +# When broadcasting (e.g. my_array[1:5] .= array_like_block), we want to treat +# ArrayLikeBlocks as scalars. +Base.broadcastable(o::ArrayLikeBlock) = Ref(o) + +_blocktype(::Type{ArrayLikeBlock{T}}) where {T} = T + +""" + GrowableArray{T,N} + +A wrapper around an `Array{T,N}`. This represents an array whose shape is not actually +known, but is presumed by the VNT to be at least as large as the largest index used to set +values into it. It is possible to call setindex!! on a GrowableArray with indices that are +out of bounds; in such a case, the GrowableArray will automatically resize itself to +accommodate the new index. +""" +struct GrowableArray{T,N} <: AbstractArray{T,N} + data::Array{T,N} +end +# I don't think this is the full AbstractArray interface, but it's enough for us (probably a +# bit too much even). +Base.size(ga::GrowableArray) = size(ga.data) +Base.axes(ga::GrowableArray) = axes(ga.data) +Base.view(ga::GrowableArray, ix...) = view(ga.data, ix...) +Base.copy(ga::GrowableArray) = GrowableArray(copy(ga.data)) +Base.:(==)(ga1::GrowableArray, ga2::GrowableArray) = ga1.data == ga2.data +Base.isequal(ga1::GrowableArray, ga2::GrowableArray) = isequal(ga1.data, ga2.data) +Base.hash(ga::GrowableArray, h::UInt) = hash(GrowableArray, hash(ga.data, h)) +Base.collect(ga::GrowableArray) = collect(ga.data) +Base.similar(ga::GrowableArray, ::Type{T}) where {T} = GrowableArray(similar(ga.data, T)) +function Base.similar(ga::GrowableArray, sz::Tuple{Vararg{Union{Integer,Base.OneTo}}}) + return GrowableArray(similar(ga.data, sz)) +end +# single-element indexing +Base.getindex(ga::GrowableArray, ix::Vararg{Int}) = getindex(ga.data, ix...) +function Base.copyto!(dest::GrowableArray, src::GrowableArray, args...) + return copyto!(dest.data, src.data, args...) +end +Base.copy!(dest::GrowableArray, src::GrowableArray) = copy!(dest.data, src.data) +Base.getindex(ga::GrowableArray, ix::CartesianIndex) = getindex(ga.data, ix) +# multi-element indexing +Base.getindex(ga::GrowableArray, ix...) = GrowableArray(getindex(ga.data, ix...)) +Base.setindex!(ga::GrowableArray, value, ix...) = setindex!(ga.data, value, ix...) +# This function is exported so we can override it! +function AbstractPPL.concretize_top_level(idx::Index, val::GrowableArray) + if any(ix -> ix isa AbstractPPL.DynamicIndex, idx.ix) + _warn_growable_array_extraction() + end + return concretize_top_level(idx, val.data) +end + +function throw_kw_error() + return throw( + ArgumentError( + "Attempted to set a value with an keyword index, but no" * + " template (or an unsuitable template) was provided for the" * + " array. A proper template is needed to determine the shape" * + " and type of the array so that the indexed data can be stored" * + " correctly.", + ), + ) +end + +# Helper functions to determine the largest index from various index types. +largest_index(ix::Integer) = ix +largest_index(r::AbstractUnitRange) = last(r) +largest_index(r::AbstractVector{<:Integer}) = maximum(r) +function largest_index(x) + return throw( + ArgumentError( + "Attempted to set a value with an index of $x, but no" * + " template was provided for the array. For such an index," * + " a template is needed to determine the shape and type of the" * + " array so that the indexed data can be stored correctly.", + ), + ) +end +function get_maximum_size_from_indices(ix...; kw...) + isempty(kw) || throw_kw_error() + return tuple(map(largest_index, ix)...) +end + +# This determines the required size for setting into the indices ix. Note that ix is not +# splatted! and this function takes no keywords! For example, if ix is (3:5,1), this +# function will return (3,); but get_maximum_size_from_indices will return (5, 1). +@generated function get_required_size_from_indices(ix::Tuple) + x = Expr(:tuple) + for (i, ti) in enumerate(ix.parameters) + if ti <: AbstractVector{<:Integer} + push!(x.args, :(length(ix[$i]))) + elseif i isa Colon + error("nope") + end + end + return x +end + +""" + PartialArray{ + ElType, + num_dims, + D<:AbstractArray{ElType,num_dims}, + M<:AbstractArray{Bool,num_dims} + } + +An array-like like structure that may only have some of its elements defined. + +One can set values in a `PartialArray` either element-by-element, or with ranges like +`arr[1:3,2] = [5,10,15]`. + +When setting values over a range of indices, the value being set can be an `AbstractArray` +whose size matches the range (in which case the values are set elementwise). If the value is +some other object, it can still be stored as an `ArrayLikeBlock`. Retrieving such a +block-value must be done with the exact same range of indices, otherwise an error is thrown. +This representation lets `PartialArray` retain array slices without flattening their values. + +If the element type of a `PartialArray` is not concrete, any call to `setindex!!` will check +if, after the new value has been set, the element type can be made more concrete. If so, +a new `PartialArray` with a more concrete element type is returned. Thus the element type +of any `PartialArray` should always be as concrete as is allowed by the elements in it. + +The internal implementation of an `PartialArray` consists of two arrays: one holding the +data and the other one being a boolean mask indicating which elements are defined. These +internal arrays may need resizing when new elements are set that have index ranges larger +than the current internal arrays. To avoid resizing too often, the internal arrays are +resized in exponentially increasing steps. This means that most `setindex!!` calls are very +fast, but some may incur substantial overhead due to resizing and copying data. It also +means that the largest index set so far determines the memory usage of the `PartialArray`. +`PartialArray`s are thus well-suited when most values in it will eventually be set. If only +a few scattered values are set, a structure like `SparseArray` may be more appropriate. +""" +struct PartialArray{ + ElType,num_dims,D<:AbstractArray{ElType,num_dims},M<:AbstractArray{Bool,num_dims} +} + data::D + mask::M + + """ + PartialArray(data::AbstractArray, mask::AbstractArray{Bool}) + + Create a new `PartialArray` with the given `data` and `mask` arrays. Note that this + constructor does not copy its arguments, so you should make sure that you do not + alias them elsewhere! + """ + function PartialArray( + data::AbstractArray{ElType,N}, mask::AbstractArray{Bool,N} + ) where {ElType,N} + if size(data) != size(mask) + throw(ArgumentError("Data and mask arrays must have the same size")) + end + return new{ElType,N,typeof(data),typeof(mask)}(data, mask) + end +end + +Base.ndims(::PartialArray{ElType,num_dims}) where {ElType,num_dims} = num_dims +Base.eltype(::PartialArray{ElType}) where {ElType} = ElType + +Base.size(pa::PartialArray) = size(pa.data) +Base.isassigned(pa::PartialArray, ix...; kw...) = isassigned(pa.data, ix...; kw...) + +function Base.copy(pa::PartialArray) + # Make a shallow copy of pa, except for any VarNamedTuple elements, which we recursively + # copy. + pa_copy = PartialArray(copy(pa.data), copy(pa.mask)) + et = eltype(pa) + if ( + VarNamedTuple <: et || + et <: VarNamedTuple || + PartialArray <: et || + et <: PartialArray + ) + @inbounds for i in eachindex(pa.mask) + if pa.mask[i] + val = @inbounds pa_copy.data[i] + if val isa VarNamedTuple || val isa PartialArray + pa_copy.data[i] = copy(val) + end + end + end + end + return pa_copy +end + +function Base.:(==)(pa1::PartialArray, pa2::PartialArray) + return size(pa1) == size(pa2) && + pa1.mask == pa2.mask && + all(pa1.data[pa1.mask] .== pa2.data[pa2.mask]) +end + +function Base.isequal(pa1::PartialArray, pa2::PartialArray) + return size(pa1) == size(pa2) && + pa1.mask == pa2.mask && + all(isequal.(pa1.data[pa1.mask], pa2.data[pa2.mask])) +end + +function Base.hash(pa::PartialArray, h::UInt) + h = hash(typeof(pa.data), h) + h = hash(size(pa.data), h) + for i in eachindex(pa.mask) + @inbounds if pa.mask[i] + h = hash(i, h) + h = hash(pa.data[i], h) + end + end + return h +end + +Base.isempty(pa::PartialArray) = !any(pa.mask) +function Base.empty(pa::PartialArray) + return PartialArray(similar(pa.data), fill!(similar(pa.mask), false)) +end +function BangBang.empty!!(pa::PartialArray) + fill!(pa.mask, false) + return pa +end + +# This is a tad hacky: We use _mapreduce_recursive which requires a prefix VarName. We give +# it the non-sense @varname(_), and then strip it away with the mapping function, returning +# only the optic. +function Base.keys(pa::PartialArray) + return _mapreduce_recursive(pair -> first(pair).optic, push!, pa, @varname(_), Any[]) +end + +# Length could be defined as a special case of mapreduce, but it's harder to keep it type +# stable that way: If the element type is abstract, we end up calling _mapreduce_recursive +# on an abstract type, which makes the type of the cumulant Any. +function Base.length(pa::PartialArray) + len = 0 + @inbounds for i in eachindex(pa.mask) + if !pa.mask[i] + continue + end + val = pa.data[i] + len += val isa VarNamedTuple || val isa PartialArray ? length(val) : 1 + end + return len +end + +""" + _concretise_eltype!!(pa::PartialArray) + +Concretise the element type of a `PartialArray`. + +Returns a new `PartialArray` with the same data and mask as `pa`, but with its element type +concretised to the most specific type that can hold all currently defined elements. + +Note that this function is fundamentally type unstable if the current element type of `pa` +is not already concrete. + +The name has a `!!` not because it mutates its argument, but because the return value +aliases memory with the argument, and is thus not independent of it. +""" +function _concretise_eltype!!(pa::PartialArray) + if isconcretetype(eltype(pa)) + return pa + end + # We could use promote_type here, instead of typejoin. However, that would e.g. + # cause Ints to be converted to Float64s, since + # promote_type(Int, Float64) == Float64, which can cause problems. See + # https://github.com/TuringLang/DynamicPPL.jl/pull/1098#discussion_r2472636188. + # Base.promote_typejoin would be like typejoin, but creates Unions out of Nothing + # and Missing, rather than falling back on Any. However, it's not exported. + new_et = typejoin((typeof(pa.data[i]) for i in eachindex(pa.mask) if pa.mask[i])...) + # TODO(mhauru) Should we check as below, or rather isconcretetype(new_et)? + # In other words, does it help to be more concrete, even if we aren't fully concrete? + if new_et === eltype(pa) + # The types of the elements do not allow for concretisation. + return pa + end + new_data = similar(pa.data, new_et) + @inbounds for i in eachindex(pa.mask) + if pa.mask[i] + new_data[i] = pa.data[i] + end + end + return PartialArray(new_data, pa.mask) +end + +""" + _can_get_arraylikeblock(pa_data::AbstractArray) + +Returns a Boolean indicating whether `pa_data` refers to a single `ArrayLikeBlock` that can +be retrieved. +""" +function _can_get_arraylikeblock(pa_data::AbstractArray) + # Check that all elements that we're trying to access are the same ArrayLikeBlock. + first_elem = first(pa_data) + if !(first_elem isa ArrayLikeBlock) || any(v -> v !== first_elem, pa_data) + return false + end + # check that there are no other elements that also refer to the same ArrayLikeBlock. + if size(pa_data) != first_elem.index_size + return false + end + # If _setindex!! works correctly, we should only be able to reach this point if all + # the elements in `val` are identical to first_elem. In this case we can return true. + return true +end + +""" + Base.getindex(pa::PartialArray, inds::Vararg{Any}; kw...) + +Obtain the value at the given indices from the `PartialArray`. This needs to be smarter than +just calling Base.getindex on the internal data array, because we need to check if the +requested indices correspond to an ArrayLikeBlock. +""" +function Base.getindex(pa::PartialArray, inds::Vararg{Any}; kw...) + # Check the mask first. We defer bounds checking to the sub-arrays. + if !(all(getindex(pa.mask, inds...; kw...))) + throw(BoundsError(pa, (inds..., kw))) + end + val = getindex(pa.data, inds...; kw...) + + # If not for ArrayLikeBlocks, at this point we could just return val directly. However, + # we need to check if val contains any ArrayLikeBlocks, and if so, make sure that that + # we are retrieving exactly that block and nothing else. + is_multiindex = _is_multiindex(pa.data, inds...; kw...) + + # The error we'll throw if the retrieval is invalid. + alb_err = ArgumentError(""" + A non-Array value set with a range of indices must be retrieved with the same + range of indices. + """) + if val isa ArrayLikeBlock + # Error if we try to get a single value, but it's an ArrayLikeBlock. + throw(alb_err) + elseif is_multiindex && + val isa AbstractArray && + (eltype(val) <: ArrayLikeBlock || ArrayLikeBlock <: eltype(val)) + # Tried to get a range of values, and at least some of them may be ArrayLikeBlocks. + # The below isempty check is deliberately kept separate from the outer elseif, + # because the outer one can be resolved at compile time. + if isempty(val) + # We need to return an empty array, but for type stability, we want to unwrap + # any ArrayLikeBlock types in the element type. + return if eltype(val) <: ArrayLikeBlock + Array{_blocktype(eltype(val)),ndims(val)}() + else + val + end + end + if _can_get_arraylikeblock(val) + return first(val).block + else + throw(alb_err) + end + elseif val isa GrowableArray + # This code path is hit for things like `vnt[@varname(x[:])]` where `x` is a PA that + # stores a GrowableArray. We warn the user that the result may be wrong. + # TODO(penelopeysm): The DynamicIndex check actually doesn't work here because if + # you index with `x[end]` for example, `inds` is already concretised outside of this + # function. + if any(ind -> ind isa AbstractPPL.DynamicIndex || ind isa Colon, inds) + _warn_growable_array_extraction() + end + return unwrap_internal_array(val) + else + return val + end +end + +function Base.haskey(pa::PartialArray, inds::Vararg{Any}; kw...) + hasall = + checkbounds(Bool, pa.mask, inds...; kw...) && all(view(pa.mask, inds...; kw...)) + + # If not for ArrayLikeBlocks, we could just return hasall directly. However, we need to + # check that if any ArrayLikeBlocks are included, they are fully included. + et = eltype(pa) + if !(et <: ArrayLikeBlock || ArrayLikeBlock <: et) + # pa can't possibly hold any ArrayLikeBlocks, so nothing to do. + return hasall + end + + if !hasall + return false + end + # From this point on we can assume that all the requested elements are set, and the only + # thing to check is that we are not partially indexing into any ArrayLikeBlocks. + subview = view(pa.data, inds...; kw...) + return if ndims(subview) == 0 + # Only one index being accessed -- just check that it isn't an ArrayLikeBlock. + isassigned(subview) && !(subview[] isa ArrayLikeBlock) + else + # Multiple indices being accessed. We need to check that we are accessing an + # ArrayLikeBlock in its entirety. + _can_get_arraylikeblock(subview) + end +end + +function BangBang.delete!!(pa::PartialArray, inds::Vararg{Any}; kw...) + fill!(view(pa.mask, inds...; kw...), false) + return pa +end + +""" + _remove_partial_blocks!!(pa::PartialArray, inds::Vararg{Any}; kw...) + +Remove any ArrayLikeBlocks that overlap with the given indices from the PartialArray. + +Note that this removes the whole block, even the parts that are within `inds`, to avoid +partially indexing into ArrayLikeBlocks. +""" +function _remove_partial_blocks!!( + pa_data::AbstractArray, pa_mask::AbstractArray{Bool}, inds::Vararg{Any}; kw... +) + et = eltype(pa_data) + if !(et <: ArrayLikeBlock || ArrayLikeBlock <: et) + # pa can't possibly hold any ArrayLikeBlocks, so nothing to do. + return pa_data, pa_mask + end + + # Generate two views, which ensures that the indices will line up + dataview = view(pa_data, inds...; kw...) + maskview = view(pa_mask, inds...; kw...) + for i in eachindex(dataview) + if maskview[i] && (dataview[i] isa ArrayLikeBlock) + val = dataview[i] + fill!(view(pa_mask, val.ix...; val.kw...), false) + end + end + return pa_data, pa_mask +end + +@inline function _ndims(f::AbstractArray, ix...; kw...) + return ndims(view(f, ix...; kw...)) +end +@inline function _ndims(f::PartialArray, ix...; kw...) + return ndims(view(f.data, ix...; kw...)) +end +@inline function _ndims(::Any, ix...; kw...) + isempty(kw) || throw_kw_error() + return _ndims_static(ix) +end +@generated function _ndims_static(::T) where {T<:Tuple} + i = 0 + for x in T.parameters + if x <: AbstractVector{<:Int} || x <: Colon + i += 1 + end + end + return :(return $i) +end +@inline function _is_multiindex(f::Union{AbstractArray,PartialArray}, ix...; kw...) + return _ndims(f, ix...; kw...) > 0 +end +@inline function _is_multiindex(::Any, ix...; kw...) + isempty(kw) || throw_kw_error() + return _is_multiindex_static(ix) +end +@generated function _is_multiindex_static(::T) where {T<:Tuple} + for x in T.parameters + if x <: AbstractVector{<:Int} || x <: Colon + return :(return true) + end + end + return :(return false) +end + +""" + _needs_arraylikeblock(pa_data::AbstractArray, value, inds::Vararg{Any}; kw...) + +Check if the given value needs to be wrapped in an `ArrayLikeBlock` when being set at the +indices `inds` in the `PartialArray` with data array `pa_data`. + +The value only depends on the types of the arguments, and should be constant propagated. +""" +function _needs_arraylikeblock(pa_data::AbstractArray, value, inds::Vararg{Any}; kw...) + return !isa(value, AbstractArray) && + !isa(value, PartialArray) && + _is_multiindex(pa_data, inds...; kw...) +end + +function grow_to_indices!!( + pa::PartialArray{T,ndims,A}, inds::Vararg{Any,ndims}; kw... +) where {T,ndims,A<:GrowableArray} + # pa.data = GrowableArray{T} + # pa.mask = GrowableArray{Bool} + required_size = get_maximum_size_from_indices(inds...; kw...) + current_size = size(pa.data) + return if any(required_size .> current_size) + old_data = pa.data.data + old_mask = pa.mask.data + # Need to make a new Array that is big enough to hold the new index. Note that these + # are just Arrays, not GrowableArrays. + new_size = map(max, required_size, current_size) + new_data = similar(old_data, new_size) + new_mask = fill(false, new_size) + # Copy the old mask into the new array. copyto! works if we are extending only the + # last dimension of the array, but if we are extending any other dimension, it will + # copy elements wrongly + for c in CartesianIndices(old_mask) + # If the old value was masked, it might be undefined. Don't copy it over. + if old_mask[c] + new_mask[c] = old_mask[c] + new_data[c] = old_data[c] + end + end + PartialArray(GrowableArray(new_data), GrowableArray(new_mask)) + else + # Can reuse the old one. + pa + end +end +function grow_to_indices!!( + pa::PartialArray{T,ndims,A}, inds::Vararg{Any,ndims2}; kw... +) where {T,ndims,ndims2,A<:GrowableArray} + return throw( + ArgumentError( + "Cannot expand a GrowableArray with $ndims dimensions" * + " using an index with $ndims2 dimensions. GrowableArrays" * + " are created when no template is provided when setting" * + " VarNames with indices, e.g. `@varname(x[1])`, which causes" * + " the VarNamedTuple to have no knowledge of the shape of `x`. To" * + " fix this, you should provide a template for `x` when creating" * + " the VarNamedTuple. Alternatively, if `x` is an array with" * + " `N` dimensions, you should always index into it with `N`" * + " indices. This means avoiding using, for example, a mixture of" * + " linear and Cartesian indexing for arrays with more than one" * + " dimension.", + ), + ) +end +grow_to_indices!!(pa::PartialArray, inds::Vararg{Any}; kw...) = pa + +function BangBang.setindex!!(pa::PartialArray, value, inds::Vararg{Any}; kw...) + # If pa.data and pa.mask are GrowableArrays, we may need to resize them before doing + # anything else. For other AbstractArrays, grow_to_indices is a no-op. + new_pa = grow_to_indices!!(pa, inds...; kw...) + new_data, new_mask = new_pa.data, new_pa.mask + + # Then delete any overlapping ArrayLikeBlocks + new_data, new_mask = _remove_partial_blocks!!(new_data, new_mask, inds...; kw...) + + if _needs_arraylikeblock(new_data, value, inds...; kw...) + idx_sz = size(@view new_data[inds..., kw...]) + alb = ArrayLikeBlock(value, inds, NamedTuple(kw), idx_sz) + new_data = setindex!!(new_data, fill(alb, idx_sz...), inds...; kw...) + fill!(view(new_mask, inds...; kw...), true) + else + if value isa PartialArray + if _is_multiindex(new_data, inds...; kw...) + # This branch occurs if we are overwriting a slice of a PA with data from + # another PA. + # + # A naive implementation of this would be: + # + # new_data = setindex!!(new_data, value.data, inds...; kw...) + # setindex!(new_mask, value.mask, inds...; kw...) + # + # However, there are two problems with this. + # + # 1. We only really want to overwrite the parts of `value` that are + # unmasked. Right now we are just copying over everything. + # 2. If it was just the above, that would be fine because it would just + # be a performance loss. The problem is that `value.data` may be + # uninitialised in the places where it is masked. Attempting to copy + # over that data would lead to an error. + # + # So we need to do it manually. :( + new_eltype = promote_type(eltype(value.data), eltype(new_data)) + new_data = if new_eltype <: eltype(new_data) + new_data + else + broadened = similar(new_data, new_eltype) + copy!(broadened, new_data) + broadened + end + new_data_view = view(new_data, inds...; kw...) + new_mask_view = view(new_mask, inds...; kw...) + for i in eachindex(new_mask_view) + if value.mask[i] + # If value.mask[i] is true, we can guarantee that value.data[i] is + # initialised to a concrete value + new_data_view[i] = value.data[i] + new_mask_view[i] = true + end + end + else + # Overwriting one element of a PA with another PA. The PA is the value + # itself! -- i.e. nested PAs! This can happen with things like x[1][1] + new_data = setindex!!(new_data, value, inds...; kw...) + setindex!(new_mask, true, inds...; kw...) + end + else + new_data = setindex!!(new_data, value, inds...; kw...) + fill!(view(new_mask, inds...; kw...), true) + end + end + + return _concretise_eltype!!(PartialArray(new_data, new_mask)) +end + +function _subset_partialarray(pa::PartialArray, inds::Vararg{Any}; kw...) + if pa.data isa GrowableArray && + any(ind -> ind isa AbstractPPL.DynamicIndex || ind isa Colon, inds) + _warn_growable_array_extraction() + end + new_data = view(pa.data, inds...; kw...) + new_mask = view(pa.mask, inds...; kw...) + return PartialArray(new_data, new_mask) +end + +Base.merge(x1::PartialArray, x2::PartialArray) = _merge(x1, x2, Val(true)) +function _merge(pa1::PartialArray, pa2::PartialArray, recurse::Val) + # If both `pa1` and `pa2` are GrowableArrays, we can grow them before merging + if pa1.data isa GrowableArray && pa2.data isa GrowableArray && ndims(pa1) == ndims(pa2) + size1 = size(pa1.data) + size2 = size(pa2.data) + new_size = map(max, size1, size2) + pa1 = grow_to_indices!!(pa1, new_size...) + pa2 = grow_to_indices!!(pa2, new_size...) + end + + # TODO(penelopeysm) In general, we should like to catch more cases (e.g. where the + # underlying Array type is different). However I don't know how to reliably check for + # that. + if axes(pa1.data) != axes(pa2.data) + throw(ArgumentError("Cannot merge PartialArrays with different axes")) + end + result = copy(pa2) + new_data, new_mask = result.data, result.mask + for i in eachindex(pa1.mask) + if pa1.mask[i] + new_elem, new_mask_val = _merge_element(pa1, pa2, i, recurse) + new_mask[i] = new_mask_val + if new_mask_val + new_data = setindex!!(new_data, new_elem, i) + end + end + end + return _concretise_eltype!!(PartialArray(new_data, new_mask)) +end + +""" + _merge_element(x1::PartialArray, x2::PartialArray, ind, ::Val{true}) + +Performs an elementwise merge of two `PartialArray`s. If both `x1` and `x2` have the element +at `ind` set, this function will attempt to merge the values recursively. + +Returns the value to be set at `ind`, plus the value of the mask at `ind`. +""" +function _merge_element(x1::PartialArray, x2::PartialArray, ind, ::Val{true}) + m1 = x1.mask[ind] + m2 = x2.mask[ind] + return if m1 && m2 + _merge(x1.data[ind], x2.data[ind], Val(true)), true + else + _merge_element(x1, x2, ind, Val(false)) + end +end + +""" + _merge_element(x1::PartialArray, x2::PartialArray, ind, ::Val{false}) + +Performs an elementwise merge of two `PartialArray`s, but does not attempt to recurse into +nested values to merge those as well. In particular, if both `x1` and `x2` have the element +at `ind` set, the value from `x2` is taken directly. + +Returns the value to be set at `ind`, plus the value of the mask at `ind`. +""" +function _merge_element(x1::PartialArray, x2::PartialArray, ind, ::Val{false}) + m1 = x1.mask[ind] + m2 = x2.mask[ind] + return if m1 && !m2 + # This is the only potential case where we need to return something in x1. However, + # there is one additional check: we only want to copy ALBs over if all the target + # indices are completely 'free' in pa2. If pa2 has any overlapping indices set, we + # shouldn't copy the ALB over. + # + # It's safe to check x1.data[ind] here because the mask was set, so it can't be + # an uninitialized value. + d1 = x1.data[ind] + if d1 isa ArrayLikeBlock + if any(view(x2.mask, d1.ix...; d1.kw...)) + # There is some overlap. Return d1 as the data (because x2.data[ind] might + # be uninitialized data!), but indicate that the mask should be false (which + # also causes setindex!! to be skipped). + d1, false + else + # No overlap -- safe to copy. Since all the target indices are free, this + # means that we will naturally copy over all the bits and pieces of the ALB + # eventually. + d1, true + end + else + # Just some other scalar value that we can copy over. + d1, true + end + else + x2.data[ind], true + end +end + +""" + unwrap_internal_array(x) + +`x` here is some data obtained by indexing with a slice into an AbstractArray or +PartialArray stored in a VNT. (It might be an AbstractArray if for example the VarName `x` +is just set to an array, rather than its individual elements being set one by one.) + +This function unwraps any PartialArray or GrowableArray wrappers, returning the underlying +Base.Array. It errors if any elements are not set, or if any elements are still +ArrayLikeBlocks. + +Sometimes, indexing with a slice into a PartialArray returns an unwrapped ArrayLikeBlock. +In that case, this function does nothing. +""" +unwrap_internal_array(x) = x +function unwrap_internal_array(pa::PartialArray) + if !(all(pa.mask)) + throw( + ArgumentError( + "Cannot extract data from PartialArray because some elements are not set." + ), + ) + end + + retval = pa.data + if retval isa GrowableArray + _warn_growable_array_extraction() + end + if eltype(retval) <: ArrayLikeBlock || ArrayLikeBlock <: eltype(retval) + for ind in eachindex(retval) + @inbounds if retval[ind] isa ArrayLikeBlock + throw( + ArgumentError( + "Cannot extract data from PartialArray when some elements " * + "are set as ArrayLikeBlocks.", + ), + ) + end + end + end + return unwrap_internal_array(retval) +end +unwrap_internal_array(ga::GrowableArray) = ga.data + +function _warn_growable_array_extraction() + @warn ( + "Returning a `Base.Array` with a presumed size based on the indices" * + " used to set values; but this may not be the actual type or size" * + " of the actual `AbstractArray` represented by the `VarNamedTuple`." * + " You should inspect the returned result to make sure that it has the" * + " correct value.\n\n" * + "To find out how to avoid this warning, please see: " * + "https://turinglang.org/docs/uri/growablearray" + ) +end +function _warn_growable_array_creation(size) + @warn ( + "Creating a growable `Base.Array` of dimension $(length(size)) to store" * + " values. This may not match the actual type or size of the actual" * + " `AbstractArray` represented by the `VarNamedTuple`.\n\n" * + " If this is not the type or size that you expect, please see:" * + " https://turinglang.org/docs/uri/growablearray" + ) +end diff --git a/src/varnamedtuple/show.jl b/src/varnamedtuple/show.jl new file mode 100644 index 00000000..eb2da640 --- /dev/null +++ b/src/varnamedtuple/show.jl @@ -0,0 +1,118 @@ +# Pretty-printing for VarNamedTuple and PartialArray. + +const ELIDED_COLOR = :light_black + +function Base.show(io::IO, vnt::VarNamedTuple) + if isempty(vnt.data) + return print(io, "VarNamedTuple()") + end + print(io, "VarNamedTuple") + show(io, vnt.data) + return nothing +end + +const MAX_KEYS_OR_INDICES = 8 + +function vnt_pretty_print(io::IO, pa::PartialArray, prefix::String, depth::Int) + # an incomplete PA is dimmed to indicate that it has missing entries. + size_style = if all(pa.mask) + (;) + else + (; color=ELIDED_COLOR) + end + print(io, "PartialArray ") + printstyled(io, "size=" * string(size(pa.data)); size_style...) + println(io, " data::" * string(typeof(pa.data))) + # can't use `keys(pa)` here as that recurses into each element. + nindices = count(pa.mask) + truncate = nindices > MAX_KEYS_OR_INDICES + active_indices = CartesianIndices(pa.data)[pa.mask] + for (i, idx) in enumerate(active_indices) + tree_symbol = i == nindices ? "└─ " : "├─ " + if truncate && i > 2 && i < nindices + if i == 3 + print(io, prefix * "│ ") + printstyled( + io, + "⋮ (" * string(nindices - 3) * " more set indices) "; + color=ELIDED_COLOR, + ) + println(io) + end + continue + end + key_name = string(Tuple(idx)) + print(io, prefix * tree_symbol) + printstyled(io, key_name; color=color_at_depth(depth)) + print(io, " => ") + v = pa.data[idx] + if v isa VarNamedTuple || v isa PartialArray + nspaces_for_key = length(key_name) + 4 + vnt_pretty_print( + io, + v, + prefix * (i == nindices ? " " : "│ ") * (" "^nspaces_for_key), + depth + 1, + ) + else + show(io, v) + end + if i < nindices + println(io) + end + end + return nothing +end +function Base.show(io::IO, ::MIME"text/plain", pa::PartialArray) + return vnt_pretty_print(io, pa, "", 0) +end + +colors = [:red, :green, :blue, :yellow, :magenta, :cyan] +color_at_depth(depth::Int) = colors[mod1(depth, length(colors))] + +function vnt_pretty_print(io::IO, vnt::VarNamedTuple, prefix::String, depth::Int) + println(io, "VarNamedTuple") + nkeys = length(keys(vnt.data)) + # If there are loads of keys, show only the first two and the last + truncate = nkeys > MAX_KEYS_OR_INDICES + for (i, (k, v)) in enumerate(zip(keys(vnt.data), values(vnt.data))) + tree_symbol = i == nkeys ? "└─ " : "├─ " + if truncate && i > 2 && i < nkeys + if i == 3 + print(io, prefix * "│ ") + printstyled( + io, "⋮ (" * string(nkeys - 3) * " more keys) "; color=ELIDED_COLOR + ) + println(io) + end + continue + end + key_name = string(k) + print(io, prefix * tree_symbol) + printstyled(io, key_name; color=color_at_depth(depth)) + print(io, " => ") + if v isa VarNamedTuple || v isa PartialArray + nspaces_for_key = length(key_name) + 4 + vnt_pretty_print( + io, + v, + prefix * (i == nkeys ? " " : "│ ") * (" "^nspaces_for_key), + depth + 1, + ) + else + show(io, v) + end + if i < nkeys + println(io) + end + end +end + +function Base.show(io::IO, ::MIME"text/plain", vnt::VarNamedTuple) + if isempty(vnt.data) + print(io, "VarNamedTuple()") + else + vnt_pretty_print(io, vnt, "", 0) + end + return nothing +end diff --git a/src/varnamedtuple/skeleton.jl b/src/varnamedtuple/skeleton.jl new file mode 100644 index 00000000..41d1508d --- /dev/null +++ b/src/varnamedtuple/skeleton.jl @@ -0,0 +1,94 @@ +""" + AbstractPPL.skeleton(vnt::VarNamedTuple) + +Create a skeleton of the given `VarNamedTuple`, which holds enough template information to +reconstruct the original `vnt` from an ordered set of (VarName, value) pairs. + +Specifically, we require that for any input `orig_vnt` that the following holds true: + +```julia +skel = skeleton(orig_vnt) +# Check roundtrip reconstruction +new_vnt = VarNamedTuple() +for (vn, val) in pairs(orig_vnt) + top_sym = AbstractPPL.getsym(vn) + template = get(skel.data, top_sym, AbstractPPL.NoTemplate()) + new_vnt = AbstractPPL.templated_setindex!!(new_vnt, val, vn, template) +end +@assert new_vnt == orig_vnt +``` + +This is accomplished by recursively traversing `vnt` and dropping any elements that do not +contain `PartialArray`s. `PartialArray`s on the other hand are replaced with arrays of the +same shape and array type, but with the elements replaced with `nothing` (since that is +space-efficient). This does not cause any issues with reconstruction type stability, since +`templated_setindex!!` does not use the array's element type when creating a new entry (it +infers the element type to use from the value that is being set). +""" +@generated function skeleton(vnt::VarNamedTuple{names,types}) where {names,types} + nms_to_include = Symbol[] + for (n, t) in zip(names, types.parameters) + if _element_needs_skeleton(t) + push!(nms_to_include, n) + end + end + return if isempty(nms_to_include) + :(VarNamedTuple()) + else + args = [:($n = skeleton(vnt.data.$n)) for n in nms_to_include] + expr = Expr(:tuple, args...) + :(VarNamedTuple($expr)) + end +end + +function _element_needs_skeleton(::Type{VarNamedTuple{names,types}}) where {names,types} + return any(_element_needs_skeleton, types.parameters) +end +_element_needs_skeleton(::Type{<:PartialArray}) = true +_element_needs_skeleton(::Type{T}) where {T} = VarNamedTuple <: T || PartialArray <: T + +# For some reason this has to be @generated, because sometimes the compiler doesn't figure +# out _element_needs_skeleton +@generated function skeleton(pa::PartialArray{T}) where {T} + if isconcretetype(T) + # We can decide from the eltype whether we need to recurse + if _element_needs_skeleton(T) + # However, constructing the actual value requires recursing on the *values* + # themselves rather than just the *type*, because the size of the array is not + # determined in the type. + # We need to be careful about this because some values in the PA may be unset. + return quote + idx = findfirst(pa.mask) + idx === nothing && error("Unexpected PartialArray with no data") + example_skeleton = skeleton(pa.data[idx]) + new_data = similar(pa.data, typeof(example_skeleton)) + fill!(new_data, example_skeleton) + return new_data + end + else + # Easy path, just fill with nothings. + return :(similar(pa.data, Nothing)) + end + else + # Need to map over each element individually, which is quite ugly. + return quote + new_data = similar(pa.data) + for i in eachindex(pa.data) + # Need to use setindex!! here because the original eltype of pa.data might + # not be broad enough to hold the skeletons themselves (which may have + # different concrete types from pa.data[i]). + if pa.mask[i] + di = pa.data[i] + if _element_needs_skeleton(typeof(di)) + BangBang.setindex!!(new_data, skeleton(di), i) + else + BangBang.setindex!!(new_data, nothing, i) + end + else + BangBang.setindex!!(new_data, nothing, i) + end + end + return new_data + end + end +end diff --git a/src/varnamedtuple/vnt.jl b/src/varnamedtuple/vnt.jl new file mode 100644 index 00000000..22fae669 --- /dev/null +++ b/src/varnamedtuple/vnt.jl @@ -0,0 +1,250 @@ +""" + VarNamedTuple{names,Values} + +A `NamedTuple`-like structure with `VarName` keys. + +`VarNamedTuple` is a data structure for storing arbitrary data, keyed by `VarName`s, in an +efficient and type stable manner. It is mainly used through `getindex`, `setindex!!`, +`templated_setindex!!`, and `haskey`, all of which only accept `VarName`s as keys. Other +notable methods are `merge` and `subset`. + +`VarNamedTuple` has an ordering to its elements, and two `VarNamedTuple`s with the same keys +and values but in different orders are considered different for equality and hashing. +Iterations such as `keys` and `values` respect this ordering. The ordering is dependent on +the order in which elements were inserted into the `VarNamedTuple`, though isn't always +equal to it. More specifically + +* Any new keys that have a joint parent `VarName` with an existing key are inserted after + that key. For instance, if one first inserts, in order, `@varname(a.x)`, `@varname(b)`, + and `@varname(a.y)`, the resulting order will be + `(@varname(a.x), @varname(a.y), @varname(b))`. +* `Index` keys`, like `@varname(a[3])` or `@varname(b[2,3,4:5])`, are always iterated + in the same order an `Array` with the same indices would be iterated. For instance, + if one first inserts, in order, `@varname(a[2])`, `@varname(b)`, and `@varname(a[1])`, + the resulting order will be `(@varname(a[1]), @varname(a[2]), @varname(b))`. + +Otherwise insertion order is respected. + +`setindex!!` and `getindex` on `VarNamedTuple` are type stable as long as one does not store +heterogeneous data under different indices of the same symbol. That is, if either + +* one sets `a[1]` and `a[2]` to be of different types, or +* if `a[1]` and `a[2]` both exist, one sets `a[1].b` without setting `a[2].b`, + +then getting values for `a[1]` or `a[2]` will not be type stable. +""" +struct VarNamedTuple{Names,Values} + data::NamedTuple{Names,Values} + + function VarNamedTuple(data::NamedTuple{Names,Values}) where {Names,Values} + return new{Names,Values}(data) + end +end + +VarNamedTuple(; kwargs...) = VarNamedTuple((; kwargs...)) + +""" + VarNamedTuple(d) + VarNamedTuple(nt::NamedTuple) + +Create a `VarNamedTuple` from a collection or a `NamedTuple`. + +Any collection `d` is assumed to be an iterable of key-value pairs, where the keys are +`VarName`s. This could be a an `AbstractDict`, a vector of `Pair`s or `Tuple`s, etc. +Alternatively, a `NamedTuple` can be passed, in which case the keys (i.e., `Symbol`s) are +converted to `VarName`s. + +Note that `VarNamedTuple` has an ordering to its elements, and two `VarNamedTuple`s with the +same keys and values but in different orders are considered different. If `d` does not +guarantee an iteration order, then the order of the elements in the resulting +`VarNamedTuple` is undefined. +""" +function VarNamedTuple(d) + vnt = VarNamedTuple() + for (k, v) in d + vnt = setindex!!(vnt, v, k) + end + return vnt +end + +Base.:(==)(vnt1::VarNamedTuple, vnt2::VarNamedTuple) = vnt1.data == vnt2.data +Base.isequal(vnt1::VarNamedTuple, vnt2::VarNamedTuple) = isequal(vnt1.data, vnt2.data) +Base.hash(vnt::VarNamedTuple, h::UInt) = hash("vnt", hash(vnt.data, h)) + +function Base.copy(vnt::VarNamedTuple{names}) where {names} + # Make a shallow copy of vnt, except for any VarNamedTuple or PartialArray elements, + # which we recursively copy. + return VarNamedTuple( + NamedTuple{names}( + map( + x -> x isa Union{VarNamedTuple,PartialArray} ? copy(x) : x, values(vnt.data) + ), + ), + ) +end + +# PartialArrays and GrowableArrays are an implementation detail of VarNamedTuple, and should +# never be the return value of getindex. Thus, we automatically convert them to dense arrays +# if needed. +function Base.getindex(vnt::VarNamedTuple, vn::VarName) + result = _getindex_optic(vnt, vn) + return unwrap_internal_array(result) +end + +Base.haskey(vnt::VarNamedTuple, vn::VarName) = _haskey_optic(vnt, vn) + +""" + AbstractPPL.VarNamedTuples.templated_setindex!!(vnt, value, vn, template) + +Assign `value` to the location in `vnt` specified by `vn`. + +The argument `template` must be provided in order to guide the creation of `PartialArray`s, +as well as to concretise any dynamic indices in `vn`. It must be an object that has the +shape of the top-level symbol in `vn`. For example: + +```julia +vnt = VarNamedTuple() +templated_setindex!!(vnt, 10, @varname(x[1]), rand(2, 2)) +``` + +Here, `rand(2, 2)` is the template for the top-level symbol `x`, which tells `setindex!!` +that `x` should be a `PartialArray` that is backed by a matrix. + +The actual data inside `template` is not needed, and `template` is never mutated by this +call. +""" +function templated_setindex!!(vnt::VarNamedTuple, value, vn::VarName, template) + return _setindex_optic!!( + vnt, value, AbstractPPL.varname_to_optic(vn), SkipTemplate{1}(template), AllowAll() + ) +end + +""" + templated_setindex_no_overwrite!!(vnt, value, vn, template) + +This is the same as `templated_setindex!!`, but throws an error if the location specified by +`vn` already has a value in `vnt`. This is useful for ensuring that values in a `VarNamedTuple` +are not set multiple times. + +!!! note + This function is entirely internal to AbstractPPL. +""" +function templated_setindex_no_overwrite!!(vnt::VarNamedTuple, value, vn::VarName, template) + return _setindex_optic!!( + vnt, + value, + AbstractPPL.varname_to_optic(vn), + SkipTemplate{1}(template), + MustNotOverwrite(vn), + ) +end + +""" + BangBang.setindex!!(vnt::VarNamedTuple, value, vn::VarName) + +This is similar to `templated_setindex!!`, but does not take a `template` argument. In +effect the `template` passed through is `NoTemplate()`. It is often the case that the +template is not actually needed (for example, if you are setting a top-level VarName: +`@varname(a)`, or if the VarName only contains Property optics, or if the arrays already +exist and you are merely updating a value inside it). In such cases, this method will work +fine, but may throw an error if the template is actually needed. Specifically, this is +likely to happen if you set a VarName with dynamic optics or colons, since in those cases +the size and shape of the underlying arrays cannot be inferred without a template. + +For example, this will error, since it is not known what `x` should be: + +```julia +vnt = VarNamedTuple() +setindex!!(vnt, 10, @varname(x[1])) +``` +""" +function BangBang.setindex!!(vnt::VarNamedTuple, value, vn::VarName) + # NOTE(penelopeysm): I _thought_ that if you allow setindex!! to reuse whatever template + # was inside the VNT, it would help with type stability. Essentially, check if + # getsym(vn) is inside the type parameter of vnt, and if so, use vnt.data.$sym as the + # template instead of NoTemplate(). But it actually made type stability worse. I don't + # know why. + return _setindex_optic!!(vnt, value, AbstractPPL.varname_to_optic(vn), NoTemplate()) +end + +""" + _has_partial_array(::Type{VarNamedTuple{Names,Values}}) where {Names,Values} + +Check if any of the types in the `Values` tuple is or contains a `PartialArray`. + +Recurses into any sub-`VarNamedTuple`s. +""" +@generated function _has_partial_array( + ::Type{VarNamedTuple{Names,Values}} +) where {Names,Values} + for T in Values.parameters + if _has_partial_array(T) + return :(return true) + end + end + return :(return false) +end + +_has_partial_array(::Type{T}) where {T} = false +_has_partial_array(::Type{<:PartialArray}) = true + +Base.empty(::VarNamedTuple) = VarNamedTuple() + +""" + empty!!(vnt::VarNamedTuple) + +Create an empty version of `vnt` in place. + +This differs from `Base.empty` in that any `PartialArray`s contained within `vnt` are kept +but have their contents deleted, rather than being removed entirely. This means that + +1) The result has a "memory" of how many dimensions different variables had, and you cannot, + for example, set `a[1,2]` after emptying a `VarNamedTuple` that had only `a[1]` defined. +2) Memory allocations may be reduced when reusing `VarNamedTuple`s, since the internal + `PartialArray`s do not need to be reallocated from scratch. +""" +@generated function BangBang.empty!!(vnt::VarNamedTuple{Names,Values}) where {Names,Values} + if !_has_partial_array(VarNamedTuple{Names,Values}) + return :(return VarNamedTuple()) + end + # Check all the fields of the NamedTuple, and keep the ones that contain PartialArrays, + # calling empty!! on them recursively. + new_names = () + new_values = () + for (name, ValType) in zip(Names, Values.parameters) + if _has_partial_array(ValType) + new_values = (new_values..., :(BangBang.empty!!(vnt.data.$name))) + new_names = (new_names..., name) + end + end + return quote + return VarNamedTuple(NamedTuple{$new_names}(($(new_values...),))) + end +end + +@generated function Base.isempty(vnt::VarNamedTuple{Names,Values}) where {Names,Values} + if isempty(Names) + return :(return true) + end + if !_has_partial_array(VarNamedTuple{Names,Values}) + return :(return false) + end + exs = Expr[] + for (name, ValType) in zip(Names, Values.parameters) + if !_has_partial_array(ValType) + return :(return false) + end + push!(exs, quote + val = vnt.data.$name + if val isa VarNamedTuple || val isa PartialArray + if !Base.isempty(val) + return false + end + else + return false + end + end) + end + push!(exs, :(return true)) + return Expr(:block, exs...) +end diff --git a/test/Project.toml b/test/Project.toml index 122f7e4e..96a08f2b 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -3,6 +3,9 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66" +Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" +ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" @@ -16,9 +19,12 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] ADTypes = "1" -AbstractPPL = "0.15" +AbstractPPL = "0.16" Accessors = "0.1" Aqua = "0.8" +BangBang = "0.4.1" +Combinatorics = "1" +ComponentArrays = "0.15" DimensionalData = "0.29, 0.30" Distributions = "0.25" Documenter = "0.26.3, 0.27, 1" diff --git a/test/runtests.jl b/test/runtests.jl index db3e8823..17006394 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -16,6 +16,7 @@ const GROUP = get(ENV, "GROUP", "All") include("varname/hasvalue.jl") include("varname/leaves.jl") include("varname/serialize.jl") + include("varnamedtuple.jl") include("of.jl") end @@ -24,7 +25,8 @@ const GROUP = get(ENV, "GROUP", "All") DocMeta.setdocmeta!( AbstractPPL, :DocTestSetup, :(using AbstractPPL); recursive=true ) - doctest(AbstractPPL; manual=false) + doctestfilters = [r"└ @ .+:[0-9]+"] + doctest(AbstractPPL; manual=false, doctestfilters) end end end diff --git a/test/varname/hasvalue.jl b/test/varname/hasvalue.jl index bf555925..085cfc90 100644 --- a/test/varname/hasvalue.jl +++ b/test/varname/hasvalue.jl @@ -34,11 +34,11 @@ using Test @test canview(@opticof(_[1:2]), x) @test canview(@opticof(_[:]), x) @test !canview(@opticof(_[4]), x) - @test canview(@opticof(_[i=1]), x) + @test canview(@opticof(_[i = 1]), x) # For some weird reason DimData does not error on these two but just warns that # there's no index j! - @test canview(@opticof(_[j=2]), x) - @test canview(@opticof(_[i=1, j=2]), x) + @test canview(@opticof(_[j = 2]), x) + @test canview(@opticof(_[i = 1, j = 2]), x) end @testset "Dict" begin @@ -246,14 +246,14 @@ end @test getvalue(x, @varname(a[1, 2])) == x.a[1, 2] @test hasvalue(x, @varname(a[:])) @test getvalue(x, @varname(a[:])) == x.a[:] - @test canview(@opticof(_[i=1]), x.a) - @test hasvalue(x, @varname(a[i=1])) - @test getvalue(x, @varname(a[i=1])) == x.a[i=1] - @test canview(@opticof(_[i=1, j=2]), x.a) - @test hasvalue(x, @varname(a[i=1, j=2])) - @test getvalue(x, @varname(a[i=1, j=2])) == x.a[i=1, j=2] - @test hasvalue(x, @varname(a[i=DD.Not(1)])) - @test getvalue(x, @varname(a[i=DD.Not(1)])) == x.a[i=DD.Not(1)] + @test canview(@opticof(_[i = 1]), x.a) + @test hasvalue(x, @varname(a[i = 1])) + @test getvalue(x, @varname(a[i = 1])) == x.a[i = 1] + @test canview(@opticof(_[i = 1, j = 2]), x.a) + @test hasvalue(x, @varname(a[i = 1, j = 2])) + @test getvalue(x, @varname(a[i = 1, j = 2])) == x.a[i = 1, j = 2] + @test hasvalue(x, @varname(a[i = DD.Not(1)])) + @test getvalue(x, @varname(a[i = DD.Not(1)])) == x.a[i = DD.Not(1)] y = (; b=DD.DimArray(randn(2, 3), (DD.X, DD.Y))) @test hasvalue(y, @varname(b)) @@ -279,6 +279,13 @@ end @test getvalue(d, @varname(x), Normal()) == 1.0 @test hasvalue(d, @varname(y[1][1]), Normal()) @test getvalue(d, @varname(y[1][1]), Normal()) == 2.0 + + vnt = @vnt begin + x := 1.0 + end + @test hasvalue(vnt, @varname(x), Normal()) + @test getvalue(vnt, @varname(x), Normal()) == 1.0 + @test !hasvalue(vnt, @varname(y), Normal()) end @testset "multivariate + matrix" begin @@ -311,6 +318,17 @@ end ) @test !hasvalue(d, @varname(y), LKJCholesky(3, 1.0); error_on_incomplete=true) + vnt = @vnt begin + @template x = (; L=zeros(2, 2)) + x.L[1, 1] := 1.0 + x.L[2, 1] := 2.0 + x.L[2, 2] := 3.0 + end + @test hasvalue(vnt, @varname(x), LKJCholesky(2, 1.0)) + @test getvalue(vnt, @varname(x), LKJCholesky(2, 1.0)) == + Cholesky(LowerTriangular([1.0 0.0; 2.0 3.0])) + @test !hasvalue(vnt, @varname(x), LKJCholesky(3, 1.0)) + d = Dict( @varname(x.U[1, 1]) => 1.0, @varname(x.U[1, 2]) => 2.0, @@ -324,6 +342,17 @@ end d, @varname(x), LKJCholesky(3, 1.0, :U); error_on_incomplete=true ) @test !hasvalue(d, @varname(y), LKJCholesky(3, 1.0, :U); error_on_incomplete=true) + + vnt = @vnt begin + @template x = (; U=zeros(2, 2)) + x.U[1, 1] := 1.0 + x.U[1, 2] := 2.0 + x.U[2, 2] := 3.0 + end + @test hasvalue(vnt, @varname(x), LKJCholesky(2, 1.0, :U)) + @test getvalue(vnt, @varname(x), LKJCholesky(2, 1.0, :U)) == + Cholesky(UpperTriangular([1.0 2.0; 0.0 3.0])) + @test !hasvalue(vnt, @varname(x), LKJCholesky(3, 1.0, :U)) end end diff --git a/test/varnamedtuple.jl b/test/varnamedtuple.jl new file mode 100644 index 00000000..8c19dbbf --- /dev/null +++ b/test/varnamedtuple.jl @@ -0,0 +1,2232 @@ +module VarNamedTupleTests + +using Combinatorics: Combinatorics +using OrderedCollections: OrderedDict +using Test: @inferred, @test, @test_throws, @testset, @test_broken, @test_logs +using AbstractPPL: + AbstractPPL, + VarName, + VarNamedTuple, + @opticof, + @varname, + @vnt, + concretize, + prefix, + subset +using AbstractPPL.VarNamedTuples: + PartialArray, + ArrayLikeBlock, + map_pairs!!, + map_values!!, + apply!!, + densify!!, + skeleton, + templated_setindex!!, + GrowableArray, + grow_to_indices!!, + MustNotOverwrite, + MustNotOverwriteError, + NoTemplate, + templated_setindex_no_overwrite!! +using BangBang: setindex!!, empty!! +using DimensionalData: DimensionalData as DD +using InvertedIndices: InvertedIndices as II +using OffsetArrays: OffsetArrays as OA +using ComponentArrays: ComponentArrays as CA + +struct GetSetTestCase + # The VarName being set. + vn + # Its value. + val + # The structure of the top-level symbol in `vn`. Its values are unused. + template + # Sub-optics of `vn` that can be accessed once it has been set in a VNT. + suboptics +end + +function test_get_set( + c::GetSetTestCase; templated_unstable=false, unstable=false, skip_setindex=false +) + @testset "Templated setindex $(c.vn)" begin + vnt = VarNamedTuple() + vnt = if templated_unstable + vnt_copy = deepcopy(vnt) + @test_throws ErrorException @inferred( + AbstractPPL.templated_setindex!!(vnt_copy, c.val, c.vn, c.template) + ) + AbstractPPL.templated_setindex!!(vnt, c.val, c.vn, c.template) + else + @inferred(AbstractPPL.templated_setindex!!(vnt, c.val, c.vn, c.template)) + end + @test @inferred(getindex(vnt, c.vn)) == c.val + for optic in c.suboptics + new_vn = AbstractPPL.append_optic(c.vn, optic) + new_val = optic(c.val) + @test @inferred(getindex(vnt, new_vn)) == new_val + end + test_invariants(vnt) + end + if !skip_setindex + @testset "setindex $(c.vn)" begin + vnt = VarNamedTuple() + vnt = if unstable + vnt_copy = deepcopy(vnt) + @test_throws ErrorException (@inferred(setindex!!(vnt_copy, c.val, c.vn))) + setindex!!(vnt, c.val, c.vn) + else + @inferred(setindex!!(vnt, c.val, c.vn)) + end + @test @inferred(getindex(vnt, c.vn)) == c.val + for optic in c.suboptics + new_vn = AbstractPPL.append_optic(c.vn, optic) + new_val = optic(c.val) + @test @inferred(getindex(vnt, new_vn)) == new_val + end + test_invariants(vnt) + end + end +end + +""" + test_invariants(vnt::VarNamedTuple; skip=()) + +Test properties that should hold for all VarNamedTuples. + +Uses @test for all the tests. Intended to be called inside a @testset. + +`skip` is a tuple of symbols indicating which tests are to be skipped. +""" +function test_invariants(vnt::VarNamedTuple) + # These will be needed repeatedly. + vnt_keys = keys(vnt) + vnt_values = values(vnt) + + # Check that for all keys in vnt, haskey is true, and resetting the value is a no-op. + for k in vnt_keys + @test haskey(vnt, k) + v = getindex(vnt, k) + # ArrayLikeBlocks and PartialArrays are implementation details, and should not be + # exposed through getindex. + @test !(v isa ArrayLikeBlock) + @test !(v isa PartialArray) + vnt2 = setindex!!(copy(vnt), v, k) + equality = (vnt == vnt2) + # The value may be `missing` if vnt itself has values that are missing. + @test equality === true || equality === missing + @test isequal(vnt, vnt2) + @test hash(vnt) == hash(vnt2) + end + + # Check that merge with an empty VarNamedTuple is a no-op. + @test isequal(merge(vnt, VarNamedTuple()), vnt) + @test isequal(merge(VarNamedTuple(), vnt), vnt) + + # Check that the VNT can be constructed back from its keys and values. + vnt4 = VarNamedTuple() + for (k, v) in zip(vnt_keys, vnt_values) + vnt4 = templated_setindex!!(vnt4, v, k, vnt.data[AbstractPPL.getsym(k)]) + end + @test isequal(vnt, vnt4) + + # Check that vnt isempty only if it has no keys + was_empty = isempty(vnt) + @test isequal(was_empty, isempty(vnt_keys)) + @test isequal(was_empty, isempty(vnt_values)) + + # Check that vnt can be emptied + @test empty(vnt) === VarNamedTuple() + emptied_vnt = empty!!(copy(vnt)) + @test isempty(emptied_vnt) + @test isempty(keys(emptied_vnt)) + @test isempty(values(emptied_vnt)) + + # Check that the copy protected the original vnt from being modified. + @test isempty(vnt) == was_empty + + # Check that map is a no-op when using identity functions. + @test isequal(map_pairs!!(pair -> pair.second, copy(vnt)), vnt) + @test isequal(map_values!!(identity, copy(vnt)), vnt) + + # Check that subsetting works as expected. + @test isequal(subset(vnt, vnt_keys), vnt) + @test isequal(subset(vnt, VarName[]), VarNamedTuple()) +end + +""" A type that has a size but is not an Array. Used in ArrayLikeBlock tests.""" +struct SizedThing{T<:Tuple} + size::T +end +Base.size(st::SizedThing) = st.size + +@testset "VarNamedTuple" begin + @testset "Construction" begin + vnt1 = VarNamedTuple() + test_invariants(vnt1) + vnt1 = setindex!!(vnt1, 1.0, @varname(a)) + vnt1 = setindex!!(vnt1, [1, 2, 3], @varname(b)) + vnt1 = setindex!!(vnt1, "a", @varname(c.d.e)) + test_invariants(vnt1) + + vnt2 = VarNamedTuple(; + a=1.0, b=[1, 2, 3], c=VarNamedTuple(; d=VarNamedTuple(; e="a")) + ) + test_invariants(vnt2) + @test vnt1 == vnt2 + + vnt3 = VarNamedTuple((; + a=1.0, b=[1, 2, 3], c=VarNamedTuple((; d=VarNamedTuple((; e="a")))) + )) + test_invariants(vnt3) + @test vnt1 == vnt3 + + vnt4 = VarNamedTuple( + OrderedDict( + @varname(a) => 1.0, @varname(b) => [1, 2, 3], @varname(c.d.e) => "a" + ), + ) + test_invariants(vnt4) + @test vnt1 == vnt4 + + vnt5 = VarNamedTuple(( + (@varname(a), 1.0), (@varname(b), [1, 2, 3]), (@varname(c.d.e), "a") + )) + test_invariants(vnt5) + @test vnt1 == vnt5 + end + + @testset "individual get/set" begin + @testset "top-level" begin + test_get_set(GetSetTestCase(@varname(a), 0.0, 0.0, [])) + test_get_set( + GetSetTestCase( + @varname(b), + [1, 2, 3], + zeros(3), + [ + @opticof(_[1]), + @opticof(_[2]), + @opticof(_[3]), + @opticof(_[1:3]), + @opticof(_[:]), + @opticof(_[1:2]) + ], + ), + ) + end + + @testset "Array indices" begin + test_get_set(GetSetTestCase(@varname(c[2]), 0.42, zeros(3), [])) + # Should still be type stable even though the eltype of the template is different, since + # the eltype is taken from the value. + test_get_set(GetSetTestCase(@varname(c[2]), "a", zeros(3), [])) + test_get_set( + GetSetTestCase( + @varname(c[2:3]), + [0.42, 1.42], + zeros(3), + [@opticof(_[1]), @opticof(_[2])], + ), + ) + end + + @testset "Properties" begin + test_get_set(GetSetTestCase(@varname(d.a), 0.42, (; a=1.0), [])) + test_get_set(GetSetTestCase(@varname(d.a.b), 0.42, (; a=(; b=1.0)), [])) + # Same as above, should still be type stable. + test_get_set(GetSetTestCase(@varname(d.a.b), "a", (; a=(; b=1.0)), [])) + test_get_set(GetSetTestCase(@varname(d.a[2]), 0.42, (; a=zeros(3)), [])) + test_get_set( + GetSetTestCase(@varname(d.a[2].e), 0.42, (; a=fill((; e=3.0), 3)), []) + ) + test_get_set( + GetSetTestCase( + @varname(e.a.b), + [1.0, 2.0], + (; a=(; b=1.0)), + [@opticof(_[1]), @opticof(_[2]), @opticof(_[:]), @opticof(_[1:2])], + ), + ) + end + + @testset "Heavily nested optics" begin + # TODO(penelopeysm): This is the only combination of things I can find for which + # templated_setindex!! is not type stable. It needs to at least have property -> + # index -> property -> index. + vnt = VarNamedTuple() + a = (; b=[nothing, (; y=zeros(2))]) + test_get_set( + GetSetTestCase(@varname(a.b[2].y[1]), 1.0, a, []); templated_unstable=true + ) + # Apart from the fact that it's 4 levels deep, to trigger the issue, the array + # template provided (in this case, the element `b`) must also have an abstract + # eltype. That's why we have to put `nothing` as the first element, rather than + # just using fill((; y=zeros(2)), 2) as the template, because that does pass. + # See for yourself: + a2 = (; b=fill((; y=zeros(2)), 2)) + test_get_set(GetSetTestCase(@varname(a.b[2].y[1]), 1.0, a2, [])) + end + + @testset "Matrices" begin + test_get_set(GetSetTestCase(@varname(f[1]), 1.0, zeros(2, 2), [])) + # Without a template this will fail because we don't know how big the second dim is, + # hence skipping the untemplated setindex. + test_get_set( + GetSetTestCase(@varname(f[1, :]), [1.0, 2.0], zeros(2, 2), []); + skip_setindex=true, + ) + test_get_set(GetSetTestCase(@varname(x[1, 1:3]), rand(3), rand(3, 3), [])) + test_get_set(GetSetTestCase(@varname(f[1, 1]), 1.0, zeros(2, 2), [])) + test_get_set(GetSetTestCase(@varname(f[1, 1]), "a", zeros(2, 2), [])) + test_get_set( + GetSetTestCase(@varname(f[2, 2].b), 1.0, fill((; b=2.0), 2, 2), []) + ) + end + + @testset "Nested single-index" begin + test_get_set(GetSetTestCase(@varname(g[2][2]), 1.0, fill(zeros(2), 2), [])) + gelem = (; a=fill(zeros(2), 2)) + g = fill(fill(gelem, 2), 2) + test_get_set(GetSetTestCase(@varname(g[2][2].a[1][1]), 1.0, g, [])) + end + + @testset "Nested multi-index" begin + test_get_set(GetSetTestCase(@varname(g[1:2][2]), 1.0, zeros(2), [])) + test_get_set( + GetSetTestCase(@varname(g[1:2][2].a), 1.0, fill((; a=1.0), 2), []); + ) + end + + @testset "OffsetArray" begin + oa = OA.OffsetArray(zeros(5), 11:15) + test_get_set(GetSetTestCase(@varname(oa[11]), 1.0, oa, []); skip_setindex=true) + test_get_set( + GetSetTestCase(@varname(x.oa[11]), 1.0, (; oa=oa), []); skip_setindex=true + ) + oa2 = OA.OffsetArray(fill((; g="a"), 5, 4), 11:15, -2:1) + test_get_set( + GetSetTestCase(@varname(oa2[11, -1].g), 1.0, oa2, []); skip_setindex=true + ) + test_get_set( + GetSetTestCase(@varname(oa2[11, -1]), 1.0, oa2, []); skip_setindex=true + ) + end + + @testset "ComponentArray" begin + ca = CA.ComponentArray(; a=1.0, b=2.0) + test_get_set(GetSetTestCase(@varname(x[1]), 1.0, ca, [])) + test_get_set(GetSetTestCase(@varname(x[2]), 2.0, ca, [])) + test_get_set(GetSetTestCase(@varname(x.a), 1.0, ca, [])) + test_get_set(GetSetTestCase(@varname(x.b), 2.0, ca, [])) + test_get_set(GetSetTestCase(@varname(x[1:2]), [1.0, 2.0], ca, [])) + + # ComponentVector with array-valued fields + ca3 = CA.ComponentArray(; a=[1.0, 2.0], b=[3.0, 4.0]) + test_get_set(GetSetTestCase(@varname(x.a), [1.0, 2.0], ca3, [])) + test_get_set(GetSetTestCase(@varname(x.b), [3.0, 4.0], ca3, [])) + test_get_set(GetSetTestCase(@varname(x.a[1]), 1.0, ca3, [])) + + # with nested fields + ca4 = CA.ComponentArray(; a=(; x=1.0, y=2.0)) + test_get_set(GetSetTestCase(@varname(x.a.x), 10.0, ca4, [])) + test_get_set(GetSetTestCase(@varname(x.a.y), 20.0, ca4, [])) + test_get_set(GetSetTestCase(@varname(x[1]), 10.0, ca4, [])) + test_get_set(GetSetTestCase(@varname(x[2]), 20.0, ca4, [])) + + # Mixed index/property access + val = rand() + vns = (@varname(x[1]), @varname(x.a)) + for set_vn in vns + vnt = AbstractPPL.templated_setindex!!(VarNamedTuple(), val, set_vn, ca) + for get_vn in vns + @test vnt[get_vn] == val + end + end + + # Check that setting one and overwriting with the other works + val = rand() + new_val = val + 1 + for (vn1, vn2) in + ((@varname(x[1]), @varname(x.a)), (@varname(x.a), @varname(x[1]))) + vnt = VarNamedTuple() + vnt = AbstractPPL.templated_setindex!!(vnt, val, vn1, ca) + @test vnt[vn1] == vnt[vn2] == val # Sanity check. + vnt = AbstractPPL.templated_setindex!!(vnt, new_val, vn2, ca) + @test vnt[vn1] == vnt[vn2] == new_val + end + + # Check that MustNotOverwrite is respected. + for vn1 in vns + vnt = AbstractPPL.templated_setindex!!(VarNamedTuple(), val, vn1, ca) + for vn2 in vns + @test_throws MustNotOverwriteError AbstractPPL.VarNamedTuples.templated_setindex_no_overwrite!!( + vnt, new_val, vn2, ca + ) + end + end + end + + @testset "InvertedIndices" begin + # TODO(penelopeysm): Templated setindex fails for II.Not(). I really don't know + # why but there is some failure in constant propagation when setting the mask + # using Not(3). Here is a minimiser that illustrates the difference between + # Not(3) and 1:2: + # + # pa = PartialArray(randn(3), fill(false, 3)) + # @descend BangBang.setindex!!(pa, [1.0, 2.0], Not(3)) + # @descend BangBang.setindex!!(pa, [1.0, 2.0], 1:2) + test_get_set( + GetSetTestCase(@varname(x[II.Not(3)]), randn(2), zeros(3), []); + templated_unstable=true, + skip_setindex=true, + ) + end + + @testset "DimensionalData" begin + da = DD.DimArray(randn(3), (DD.X)) + test_get_set( + GetSetTestCase(@varname(da[DD.X(2)]), 1.0, da, []); skip_setindex=true + ) + test_get_set( + GetSetTestCase(@varname(x.da[DD.X(2)]), 1.0, (; da=da), []); + skip_setindex=true, + ) + da2 = DD.DimArray(fill((; g="a"), 3, 5), (DD.X, DD.Y)) + test_get_set( + GetSetTestCase(@varname(da2[DD.X(2), DD.Y(1)]), 1.0, da2, []); + skip_setindex=true, + ) + test_get_set( + GetSetTestCase(@varname(da2[DD.X(2), DD.Y(1)].g), "b", da2, []); + skip_setindex=true, + ) + end + end + + @testset "KeyError for missing properties" begin + vnt = @vnt begin + x.a := 1.0 + end + # Should throw KeyError for missing top-level symbol + @test_throws KeyError vnt[@varname(y)] + # Should throw KeyError for missing nested property + @test_throws KeyError vnt[@varname(x.b)] + # Sanity check: accessing existing property should work + @test vnt[@varname(x.a)] == 1.0 + end + + @testset "haskey on PartialArray" begin + @testset "no ALBs" begin + vnt = @vnt begin + @template x = zeros(2) + x[1] := 1.0 + x[2] := 2.0 + end + # Sanity check + @test haskey(vnt, @varname(x[1])) + @test haskey(vnt, @varname(x[2])) + @test !haskey(vnt, @varname(x[3])) + # These should also be true + @test haskey(vnt, @varname(x[1:2])) + @test haskey(vnt, @varname(x[1:1])) + @test haskey(vnt, @varname(x[2:2])) + @test haskey(vnt, @varname(x[:])) + end + + @testset "mixed values + ALBs" begin + vnt2 = @vnt begin + @template x = zeros(3) + x[1] := 1.0 + x[2:3] := SizedThing((2,)) + end + @test haskey(vnt2, @varname(x[1])) + @test haskey(vnt2, @varname(x[2:3])) + @test !haskey(vnt2, @varname(x[:])) + @test !haskey(vnt2, @varname(x[1:2])) + @test !haskey(vnt2, @varname(x)) + end + + @testset "only ALBs" begin + vnt3 = @vnt begin + @template x = zeros(3) + x[1:3] := SizedThing((3,)) + end + @test haskey(vnt3, @varname(x[1:3])) + @test haskey(vnt3, @varname(x[:])) + @test haskey(vnt3, @varname(x)) + @test !haskey(vnt3, @varname(x[1:2])) + @test !haskey(vnt3, @varname(x[1])) + end + end + + @testset "Setting to same variable multiple times" begin + # TODO(penelopeysm) write these + end + + @testset "Chain of sets and gets" begin + vnt = VarNamedTuple() + vnt = @inferred(setindex!!(vnt, 32.0, @varname(a))) + @test @inferred(getindex(vnt, @varname(a))) == 32.0 + @test haskey(vnt, @varname(a)) + @test !haskey(vnt, @varname(b)) + test_invariants(vnt) + + vnt = @inferred(setindex!!(vnt, [1, 2, 3], @varname(b))) + @test @inferred(getindex(vnt, @varname(b))) == [1, 2, 3] + @test @inferred(getindex(vnt, @varname(b[2]))) == 2 + @test haskey(vnt, @varname(b)) + @test haskey(vnt, @varname(b[1])) + @test haskey(vnt, @varname(b[1:3])) + @test !haskey(vnt, @varname(b[4])) + test_invariants(vnt) + + vnt = @inferred(setindex!!(vnt, 64.0, @varname(a))) + @test @inferred(getindex(vnt, @varname(a))) == 64.0 + @test @inferred(getindex(vnt, @varname(b))) == [1, 2, 3] + test_invariants(vnt) + + vnt = @inferred(setindex!!(vnt, 15, @varname(b[2]))) + @test @inferred(getindex(vnt, @varname(b))) == [1, 15, 3] + @test @inferred(getindex(vnt, @varname(b[2]))) == 15 + test_invariants(vnt) + + vnt = @inferred(setindex!!(vnt, [10], @varname(c.x.y))) + @test @inferred(getindex(vnt, @varname(c.x.y))) == [10] + test_invariants(vnt) + + vnt = @inferred(setindex!!(vnt, 11, @varname(c.x.y[1]))) + @test @inferred(getindex(vnt, @varname(c.x.y))) == [11] + @test @inferred(getindex(vnt, @varname(c.x.y[1]))) == 11 + test_invariants(vnt) + + vnt = @inferred(templated_setindex!!(vnt, -1.0, @varname(d[4]), randn(4))) + @test @inferred(getindex(vnt, @varname(d[4]))) == -1.0 + test_invariants(vnt) + + vnt = @inferred(setindex!!(vnt, -2.0, @varname(d[4]))) + @test @inferred(getindex(vnt, @varname(d[4]))) == -2.0 + test_invariants(vnt) + + e = (; f=fill((; g=(; h=fill((; i=0.0), 2))), 3)) + vnt = @inferred(templated_setindex!!(vnt, 1.0, @varname(e.f[3].g.h[2].i), e)) + @test @inferred(getindex(vnt, @varname(e.f[3].g.h[2].i))) == 1.0 + @test haskey(vnt, @varname(e.f[3].g.h[2].i)) + @test !haskey(vnt, @varname(e.f[2].g.h[2].i)) + test_invariants(vnt) + + # TODO(penelopeysm) This one fails type stability. + # https://github.com/TuringLang/DynamicPPL.jl/issues/1207 + vnt = setindex!!(vnt, 2.0, @varname(e.f[3].g.h[2].i)) + @test @inferred(getindex(vnt, @varname(e.f[3].g.h[2].i))) == 2.0 + test_invariants(vnt) + + jval = fill(1.0, 4) + vnt = @inferred(templated_setindex!!(vnt, jval, @varname(j[1:4]), zeros(4))) + @test @inferred(getindex(vnt, @varname(j[1:4]))) == jval + @test @inferred(getindex(vnt, @varname(j[2]))) == jval[2] + @test haskey(vnt, @varname(j[4])) + @test !haskey(vnt, @varname(j[5])) + @test_throws BoundsError getindex(vnt, @varname(j[5])) + test_invariants(vnt) + + j2val = fill(2.0, 4) + vnt = @inferred(templated_setindex!!(vnt, j2val, @varname(j2[2:5]), zeros(5))) + @test_throws BoundsError getindex(vnt, @varname(j2[1])) + @test @inferred(getindex(vnt, @varname(j2[2:5]))) == j2val + @test haskey(vnt, @varname(j2[5])) + test_invariants(vnt) + + arr = fill(2.0, (4, 2)) + vn = @varname(k.l[2:5, 3, 1:2, 2]) + vnt = @inferred(templated_setindex!!(vnt, arr, vn, (; l=zeros(5, 3, 2, 2)))) + @test @inferred(getindex(vnt, vn)) == arr + # A subset of the elements set just now. + @test @inferred(getindex(vnt, @varname(k.l[2, 3, 1:2, 2]))) == fill(2.0, 2) + test_invariants(vnt) + + # Not enough, or too many, indices. + @test_throws BoundsError setindex!!(vnt, 0.0, @varname(k.l[1, 2, 3])) + @test_throws BoundsError setindex!!(vnt, 0.0, @varname(k.l[1, 2, 3, 4, 5])) + + vnt = @inferred(templated_setindex!!(vnt, 1.0, @varname(m[2]), randn(3))) + vnt = @inferred(setindex!!(vnt, 1.0, @varname(m[3]))) + @test @inferred(getindex(vnt, @varname(m[2:3]))) == [1.0, 1.0] + @test !haskey(vnt, @varname(m[1])) + test_invariants(vnt) + vnt = @inferred(setindex!!(vnt, 1.0, @varname(m[1]))) + @test @inferred(getindex(vnt, @varname(m[1:3]))) == ones(3) + @test @inferred(getindex(vnt, @varname(m))) == ones(3) + @test @inferred(getindex(vnt, @varname(m[:]))) == ones(3) + test_invariants(vnt) + + # The below tests are mostly significant for the type stability aspect. For the last + # test to pass, PartialArray needs to actively tighten its eltype when possible. + vnt = VarNamedTuple() + vnt = @inferred(templated_setindex!!(vnt, 1.0, @varname(n[1].a), [0.0, 0.0])) + @test @inferred(getindex(vnt, @varname(n[1].a))) == 1.0 + vnt = @inferred(setindex!!(vnt, 1.0, @varname(n[2].a))) + @test @inferred(getindex(vnt, @varname(n[2].a))) == 1.0 + # This can't be type stable, because n[1] has inhomogeneous types. + vnt = setindex!!(vnt, 1.0, @varname(n[1].b)) + @test getindex(vnt, @varname(n[1].b)) == 1.0 + # The setindex!! call can't be type stable either, but it should return a + # VarNamedTuple with a concrete element type, and hence getindex can be inferred. + vnt = setindex!!(vnt, 1.0, @varname(n[2].b)) + @test @inferred(getindex(vnt, @varname(n[2].b))) == 1.0 + test_invariants(vnt) + + # Some funky Symbols in VarNames. + vn1 = VarName{Symbol("a b c")}() + vnt = @inferred(setindex!!(vnt, 2, vn1)) + @test @inferred(getindex(vnt, vn1)) == 2 + test_invariants(vnt) + vn2 = VarName{Symbol("1")}() + vnt = @inferred(setindex!!(vnt, 3, vn2)) + @test @inferred(getindex(vnt, vn2)) == 3 + test_invariants(vnt) + vn3 = VarName{Symbol("?!")}() + vnt = @inferred(setindex!!(vnt, 4, vn3)) + @test @inferred(getindex(vnt, vn3)) == 4 + test_invariants(vnt) + vnt = VarNamedTuple() + vn4 = prefix(prefix(vn1, vn2), vn3) + vnt = @inferred(setindex!!(vnt, 5, vn4)) + @test @inferred(getindex(vnt, vn4)) == 5 + test_invariants(vnt) + vn5 = prefix(prefix(vn3, vn2), vn1) + vnt = @inferred(setindex!!(vnt, 6, vn5)) + @test @inferred(getindex(vnt, vn5)) == 6 + test_invariants(vnt) + + vnt = VarNamedTuple() + x = [1, 2, 3] + vn = @varname(x[:]) + vnt = @inferred(templated_setindex!!(vnt, x, vn, x)) + @test haskey(vnt, vn) + @test @inferred(getindex(vnt, vn)) == x + test_invariants(vnt) + + vnt = VarNamedTuple() + vn = @varname(y[:]) + vnt = @inferred(templated_setindex!!(vnt, SizedThing((3,)), vn, randn(3))) + @test haskey(vnt, vn) + @test vn in keys(vnt) + @test @inferred(getindex(vnt, vn)) == SizedThing((3,)) + test_invariants(vnt) + + vnt = VarNamedTuple() + x = fill("a", (3, 4)) + a = (; b=fill((; c=(; d=zeros(1, 5, 2, 4, 1))), 3)) + vn = @varname(a.b[3].c.d[1, 3:5, 2, :, 1]) + vnt = @inferred(templated_setindex!!(vnt, x, vn, a)) + @test haskey(vnt, vn) + @test @inferred(getindex(vnt, vn)) == x + test_invariants(vnt) + + # Indices on indices + vnt = VarNamedTuple() + vnt = @inferred(templated_setindex!!(vnt, 1, @varname(a[1][1]), [[randn()]])) + @test @inferred(getindex(vnt, @varname(a[1][1]))) == 1 + vnt = @inferred(templated_setindex!!(vnt, 1, @varname(ab[1:2][1]), randn(2))) + @test @inferred(getindex(vnt, @varname(ab[1]))) == 1 + @test @inferred(getindex(vnt, @varname(ab[1:2][1]))) == 1 + @test @inferred(getindex(vnt, @varname(ab[:][1]))) == 1 + @test_throws BoundsError getindex(vnt, @varname(ab[2])) + vnt = @inferred( + templated_setindex!!(vnt, [1], @varname(b[1].c[1]), [(; c=randn(1))]) + ) + @test @inferred(getindex(vnt, @varname(b[1].c[1]))) == [1] + # TODO(penelopeysm) These all have to be changed to templated_setindex!!, but I + # can't be bothered right now. + # vnt = @inferred(setindex!!(vnt, [1], @varname(e[3, 2].f[2, 2][10, 10]))) + # @test @inferred(getindex(vnt, @varname(e[3, 2].f[2, 2][10, 10]))) == [1] + # vnt = @inferred(setindex!!(vnt, [1], @varname(g[3, 2][10, 10].h[2, 2]))) + # @test @inferred(getindex(vnt, @varname(g[3, 2][10, 10].h[2, 2]))) == [1] + end + + @testset "GrowableArray behaviour" begin + @testset "similar" begin + ga = GrowableArray(ones(2, 3)) + @test similar(ga, (3, 4)) isa GrowableArray + @test similar(ga, axes(ga)) isa GrowableArray + end + + @testset "grow_to_indices!! on PartialArray->GrowableArray" begin + pa = PartialArray(GrowableArray(ones(2)), GrowableArray([false, true])) + @test_throws BoundsError pa[1] + @test pa[2] == 1.0 + pa = grow_to_indices!!(pa, 1:3) + @test size(pa.data) == (3,) + @test_throws BoundsError pa[1] + @test pa[2] == 1.0 + # pa[3] should still error... + @test_throws BoundsError pa[3] + # ... but the underlying data should be there + @test pa.mask[3] == false + end + + @testset "grow_to_indices!! can't change dimensionality" begin + pa = PartialArray(GrowableArray(ones(2)), GrowableArray([false, true])) + @test_throws ArgumentError grow_to_indices!!(pa, 1, 2) + pa = PartialArray(GrowableArray(randn(2, 2)), GrowableArray(fill(false, 2, 2))) + @test_throws ArgumentError grow_to_indices!!(pa, 1) + end + + @testset "Data is correctly copied when expanding" begin + x = randn(2, 3) + vnt = VarNamedTuple() + for c in CartesianIndices((2, 3)) + vnt = setindex!!(vnt, x[c], @varname(x[c.I...])) + end + @test vnt[@varname(x)] == x + end + + @testset "GrowableArrays and slices" begin + # see e.g. https://github.com/TuringLang/DynamicPPL.jl/issues/1216 + vnt = VarNamedTuple() + vnt = @inferred(setindex!!(vnt, 1.0, @varname(x[1:2][1]))) + vnt = @inferred(setindex!!(vnt, 1.0, @varname(x[2:3][1]))) + @test vnt[@varname(x[1])] == 1.0 + @test vnt[@varname(x[2])] == 1.0 + @test_throws BoundsError vnt[@varname(x[3])] + + # check overlapping indices + vnt = VarNamedTuple() + vnt = @inferred(setindex!!(vnt, 1.0, @varname(x[1:2][2]))) + vnt = @inferred(setindex!!(vnt, 2.0, @varname(x[2:3][1]))) + @test_throws BoundsError vnt[@varname(x[1])] + @test vnt[@varname(x[2])] == 2.0 + @test_throws BoundsError vnt[@varname(x[3])] + + # slices of slices + vnt = VarNamedTuple() + vnt = @inferred(setindex!!(vnt, [1.0, 2.0], @varname(x[2:6][2:3]))) + vnt = @inferred(setindex!!(vnt, [3.0, 4.0], @varname(x[2:6][4:5]))) + @test vnt[@varname(x[3:6])] == [1.0, 2.0, 3.0, 4.0] + @test_throws BoundsError vnt[@varname(x[1])] + @test_throws BoundsError vnt[@varname(x[2])] + @test_throws BoundsError vnt[@varname(x[7])] + + # overlapping slices of slices + vnt = VarNamedTuple() + vnt = @inferred(setindex!!(vnt, [1.0, 2.0], @varname(x[2:6][2:3]))) + vnt = @inferred(setindex!!(vnt, [3.0, 4.0], @varname(x[3:5][2:3]))) + @test vnt[@varname(x[3:5])] == [1.0, 3.0, 4.0] + @test_throws BoundsError vnt[@varname(x[1])] + @test_throws BoundsError vnt[@varname(x[2])] + @test_throws BoundsError vnt[@varname(x[6])] + @test_throws BoundsError vnt[@varname(x[7])] + end + + @testset "GrowableArrays only need to grow when setting slices" begin + # For VarNames like x[1:2][1], we need to make the inner GrowableArray (which is + # initially created with length 1) grow to fit the indices 1:2. However, for + # x[1][1,2], we need to make sure that we don't try to make the inner + # GrowableArray (that has 2 dims) grow to fit index 1, which would error. + vnt = VarNamedTuple() + vnt = @inferred(setindex!!(vnt, 1.0, @varname(x[1][1, 2]))) + vnt = @inferred(setindex!!(vnt, 2.0, @varname(x[1][1, 2]))) + @test @inferred(vnt[@varname(x[1][1, 2])]) == 2.0 + end + + @testset "Merging" begin + # Note that there are a couple more merge tests with ALBs under the merge + # section. + @testset "non-overlapping" begin + v1 = setindex!!(VarNamedTuple(), 1.0, @varname(x[1])) + v2 = setindex!!(VarNamedTuple(), 2.0, @varname(x[2])) + vmerged = @inferred(merge(v1, v2)) + @test @inferred(vmerged[@varname(x[1])]) == 1.0 + @test @inferred(vmerged[@varname(x[2])]) == 2.0 + @test length(keys(vmerged)) == 2 + end + + @testset "overlapping" begin + v1 = setindex!!(VarNamedTuple(), [1.0, 2.0], @varname(x[1:2])) + v2 = setindex!!(VarNamedTuple(), [3.0, 4.0], @varname(x[2:3])) + vmerged = @inferred(merge(v1, v2)) + @test @inferred(vmerged[@varname(x[1])]) == 1.0 + @test @inferred(vmerged[@varname(x[2])]) == 3.0 + @test @inferred(vmerged[@varname(x[3])]) == 4.0 + @test length(keys(vmerged)) == 3 + end + end + + @testset "Warning when setting untemplated values" begin + warnmsg = r"Creating a growable `Base.Array`" + @test_logs (:warn, warnmsg) vnt = setindex!!( + VarNamedTuple(), [1.0, 2.0], @varname(x[1:2]) + ) + @test_logs (:warn, warnmsg) vnt = setindex!!( + VarNamedTuple(), 1.0, @varname(x[1]) + ) + @test_logs (:warn, warnmsg) vnt = setindex!!( + VarNamedTuple(), 1.0, @varname(x[2]) + ) + # This shouldn't warn, because `x` is just the whole vector + @test_logs vnt = setindex!!(VarNamedTuple(), [1.0, 2.0], @varname(x)) + # With a template it shouldn't warn + @test_logs vnt = templated_setindex!!( + VarNamedTuple(), 1.0, @varname(x[1]), zeros(2) + ) + end + + @testset "Warning when trying to extract values" begin + vnt = @vnt begin + x[1] := 1.0 + x[2] := 2.0 + end + warnmsg = r"Returning a `Base.Array` with" + # These should warn. There are various calls to _warn_growable_array() inside + # the VNT code to handle each of these cases. + @test_logs (:warn, warnmsg) vnt[@varname(x[:][1])] + @test_logs (:warn, warnmsg) vnt[@varname(x[:])] + @test_logs (:warn, warnmsg) vnt[@varname(x)] + @test_logs (:warn, warnmsg) vnt[@varname(x[begin])] + @test_logs (:warn, warnmsg) vnt[@varname(x[end])] + @test_logs (:warn, warnmsg) vnt[@varname(x[1:end])] + # These shouldn't warn because they are safe + @test_logs vnt[@varname(x[1:2])] + @test_logs vnt[@varname(x[1])] + end + end + + @testset "multiindices" begin + # These tests essentially check that: + # (1) setting a slice works as expected + # (2) setting an index of a slice works as expected + # (3) setting a slice of a slice works as expected + # + # There used to be a litany of bugs around this. See DynamicPPL #1205. + @testset for (index_chains, template_length) in [ + ((1:2, 1), 3), + ((2:3, 1), 3), + ((1:2, 2), 3), + ((2:3, 2), 3), + ((1:3, 1:2), 4), + ((1:3, 2:3), 4), + ((2:4, 1:2), 4), + ((2:4, 2:3), 4), + ((1:3, 1:2, 1), 4), + ((1:3, 1:2, 2), 4), + ((2:4, 2:3, 1), 4), + ((2:4, 2:3, 2), 4), + ((1:3, 2:3, 1), 4), + ((1:3, 2:3, 2), 4), + ((2:4, 1:2, 1), 4), + ((2:4, 1:2, 2), 4), + ((1:3, :), 4), + ((2:4, :), 4), + ((:, 1:3), 4), + ((:, 2:4), 4), + ((1:3, :, 1:2), 4), + ((2:4, :, 1:2), 4), + ((1:3, :, 2:3), 4), + ((2:4, :, 2:3), 4), + ((1:2, :, 1), 3), + ((1:2, :, 2), 3), + ((2:3, :, 1), 3), + ((2:3, :, 2), 3), + ] + # I couldn't figure out a better way to do this. `@varname(x[index_chains...])` + # works, but it gives the wrong thing: it gives you `x[1:2, 1]` instead of + # `x[1:2][1]`. + vn = if length(index_chains) == 2 + @varname(x[index_chains[1]][index_chains[2]]) + elseif length(index_chains) == 3 + @varname(x[index_chains[1]][index_chains[2]][index_chains[3]]) + else + error("unsupported length") + end + # This gets the final position where the values should end up -- + # get_collapsed_indices((1:3, 2:3, 1), nothing) = 2 + function get_collapsed_indices(chain, init) + if isnothing(init) || init isa Colon + return get_collapsed_indices(Base.tail(chain), first(chain)) + end + isempty(chain) && return init + inds = first(chain) + return if inds isa Colon + get_collapsed_indices(Base.tail(chain), init) + else + get_collapsed_indices(Base.tail(chain), init[inds]) + end + end + final_indices = get_collapsed_indices(index_chains, nothing) + + val = if final_indices isa Integer + 1.0 + else + fill(1.0, length(final_indices)) + end + + for (func!!, args) in + ((templated_setindex!!, (zeros(template_length),)), (setindex!!, ())) + vnt = @inferred( + templated_setindex!!(VarNamedTuple(), val, vn, zeros(template_length)) + ) + @test @inferred(getindex(vnt, vn)) == val + normalized_vn = @varname(x[final_indices]) + @test @inferred(getindex(vnt, normalized_vn)) == val + + # Check that other indices are not set + for i in 1:template_length + if !(i in final_indices) + @test_throws BoundsError getindex(vnt, @varname(x[i])) + end + end + end + end + end + + @testset "chained multiindices [1:2][1] -> [1:2][2]" begin + # These are particularly finicky to deal with. See + # https://github.com/TuringLang/DynamicPPL.jl/issues/1205. + # + # The aim of these tests is to make sure that when you set [1:2][2], you don't + # lose [1:2][1]. This could happen if you're not careful because when you + # set [1:2][2], you get back a length-2 PartialArray with only the second element + # set, and you could overwrite the entire [1:2] entry in the target PartialArray, + + x = zeros(2) + for (func!!, args) in ((templated_setindex!!, (x,)), (setindex!!, ())) + vnt = VarNamedTuple() + + vnt = @inferred(func!!(vnt, 1.0, @varname(x[1:2][1]), args...)) + @test @inferred(getindex(vnt, @varname(x[1:2][1]))) == 1.0 + @test @inferred(getindex(vnt, @varname(x[1]))) == 1.0 + @test_throws BoundsError getindex(vnt, @varname(x[2])) + test_invariants(vnt) + + # Now set the second index + vnt = @inferred(func!!(vnt, 2.0, @varname(x[1:2][2]), args...)) + @test @inferred(getindex(vnt, @varname(x[1:2][2]))) == 2.0 + @test @inferred(getindex(vnt, @varname(x[2]))) == 2.0 + test_invariants(vnt) + + # Check that the first index is still correct + @test @inferred(getindex(vnt, @varname(x[1]))) == 1.0 + # Check that we can get both indices + @test @inferred(getindex(vnt, @varname(x[1:2]))) == [1.0, 2.0] + @test @inferred(getindex(vnt, @varname(x[:]))) == [1.0, 2.0] + @test @inferred(getindex(vnt, @varname(x))) == [1.0, 2.0] + + @test keys(vnt) == [@varname(x[1]), @varname(x[2])] + @test values(vnt) == [1.0, 2.0] + end + end + + @testset "chained multiindices [1:2][1] -> [end:end][1]" begin + # See https://github.com/TuringLang/DynamicPPL.jl/issues/1205. + # Note that this doesn't work with untemplated setindex!! because `end` will be + # wrong. + vnt = VarNamedTuple() + x = zeros(3) + + vnt = @inferred(templated_setindex!!(vnt, 1.0, @varname(x[1:2][1]), x)) + @test @inferred(getindex(vnt, @varname(x[1:2][1]))) == 1.0 + @test @inferred(getindex(vnt, @varname(x[1]))) == 1.0 + @test_throws BoundsError getindex(vnt, @varname(x[2])) + @test_throws BoundsError getindex(vnt, @varname(x[3])) + test_invariants(vnt) + + # Now set the second index + vnt = @inferred(templated_setindex!!(vnt, 2.0, @varname(x[end:end][1]), x)) + @test @inferred(getindex(vnt, @varname(x[end:end][1]))) == 2.0 + @test @inferred(getindex(vnt, @varname(x[3]))) == 2.0 + test_invariants(vnt) + + # Check that the first index is still correct + @test @inferred(getindex(vnt, @varname(x[1]))) == 1.0 + # Check that we can get both indices + @test @inferred(getindex(vnt, @varname(x[[1, 3]]))) == [1.0, 2.0] + # Check that we can't get all three indices + @test_throws BoundsError getindex(vnt, @varname(x[1:3])) + @test_throws BoundsError getindex(vnt, @varname(x[:])) + @test_throws ArgumentError getindex(vnt, @varname(x)) + + @test keys(vnt) == [@varname(x[1]), @varname(x[3])] + @test values(vnt) == [1.0, 2.0] + end + + @testset "equality and hash" begin + # Test all combinations of having or not having the below values set, and having + # them set to any of the possible_values, and check that isequal and == return the + # expected value. + # NOTE: Be very careful adding new values to these sets. The below test has three + # nested loops over Combinatorics.combinations, the run time can explode very, very + # quickly. + b = rand(3) + c = (; d=[nothing, (; e=rand()), nothing]) + varnames_and_templates = ( + (@varname(b[1]), b), (@varname(b[3]), b), (@varname(c.d[2].e), c) + ) + possible_values = (missing, 1, -0.0, 0.0) + for vn_template_set in Combinatorics.combinations(varnames_and_templates) + valuesets1 = Combinatorics.with_replacement_combinations( + possible_values, length(vn_template_set) + ) + valuesets2 = Combinatorics.with_replacement_combinations( + possible_values, length(vn_template_set) + ) + for vset1 in valuesets1, vset2 in valuesets2 + vnt1 = VarNamedTuple() + vnt2 = VarNamedTuple() + expected_isequal = true + expected_doubleequal = true + for ((vn, template), v1, v2) in zip(vn_template_set, vset1, vset2) + vnt1 = templated_setindex!!(vnt1, v1, vn, template) + vnt2 = templated_setindex!!(vnt2, v2, vn, template) + expected_isequal = expected_isequal & isequal(v1, v2) + expected_doubleequal = expected_doubleequal & (v1 == v2) + end + test_invariants(vnt1) + test_invariants(vnt2) + @test isequal(vnt1, vnt2) == expected_isequal + @test (vnt1 == vnt2) === expected_doubleequal + if expected_isequal + @test hash(vnt1) == hash(vnt2) + end + end + end + end + + @testset "merge" begin + vnt1 = VarNamedTuple() + vnt2 = VarNamedTuple() + expected_merge = VarNamedTuple() + @test @inferred(merge(vnt1, vnt2)) == expected_merge + + vnt1 = setindex!!(vnt1, 1.0, @varname(a)) + vnt2 = setindex!!(vnt2, 2.0, @varname(b)) + vnt1 = setindex!!(vnt1, 1, @varname(c)) + vnt2 = setindex!!(vnt2, 2, @varname(c)) + expected_merge = setindex!!(expected_merge, 1.0, @varname(a)) + expected_merge = setindex!!(expected_merge, 2, @varname(c)) + expected_merge = setindex!!(expected_merge, 2.0, @varname(b)) + @test @inferred(merge(vnt1, vnt2)) == expected_merge + test_invariants(vnt1) + test_invariants(vnt2) + + vnt1 = VarNamedTuple() + vnt2 = VarNamedTuple() + expected_merge = VarNamedTuple() + vnt1 = setindex!!(vnt1, [1], @varname(d.a)) + vnt2 = setindex!!(vnt2, [2, 2], @varname(d.b)) + vnt1 = setindex!!(vnt1, [1], @varname(d.c)) + vnt2 = setindex!!(vnt2, [2, 2], @varname(d.c)) + expected_merge = setindex!!(expected_merge, [1], @varname(d.a)) + expected_merge = setindex!!(expected_merge, [2, 2], @varname(d.c)) + expected_merge = setindex!!(expected_merge, [2, 2], @varname(d.b)) + @test @inferred(merge(vnt1, vnt2)) == expected_merge + + e = (; a=rand(11), b=[rand(20) for _ in 1:4]) + vnt1 = templated_setindex!!(vnt1, 1, @varname(e.a[1]), e) + vnt2 = templated_setindex!!(vnt2, 2, @varname(e.a[2]), e) + expected_merge = templated_setindex!!(expected_merge, 1, @varname(e.a[1]), e) + expected_merge = setindex!!(expected_merge, 2, @varname(e.a[2])) + vnt1 = setindex!!(vnt1, 1, @varname(e.a[3])) + vnt2 = setindex!!(vnt2, 2, @varname(e.a[3])) + expected_merge = setindex!!(expected_merge, 2, @varname(e.a[3])) + @test @inferred(merge(vnt1, vnt2)) == expected_merge + + vnt1 = setindex!!(vnt1, fill(1, 4), @varname(e.a[7:10])) + vnt2 = setindex!!(vnt2, fill(2, 4), @varname(e.a[8:11])) + expected_merge = setindex!!(expected_merge, 1, @varname(e.a[7])) + expected_merge = setindex!!(expected_merge, fill(2, 4), @varname(e.a[8:11])) + @test @inferred(merge(vnt1, vnt2)) == expected_merge + + vnt1 = templated_setindex!!(vnt1, 1, @varname(e.b[1][13]), e) + vnt2 = templated_setindex!!(vnt2, 2, @varname(e.b[2][13]), e) + expected_merge = templated_setindex!!(expected_merge, 1, @varname(e.b[1][13]), e) + expected_merge = templated_setindex!!(expected_merge, 2, @varname(e.b[2][13]), e) + vnt1 = templated_setindex!!(vnt1, 1, @varname(e.b[3][13]), e) + vnt2 = templated_setindex!!(vnt2, 2, @varname(e.b[3][13]), e) + expected_merge = templated_setindex!!(expected_merge, 2, @varname(e.b[3][13]), e) + @test @inferred(merge(vnt1, vnt2)) == expected_merge + vnt1 = templated_setindex!!(vnt1, 1, @varname(e.b[4][13]), e) + vnt2 = templated_setindex!!(vnt2, 2, @varname(e.b[4][14]), e) + expected_merge = templated_setindex!!(expected_merge, 1, @varname(e.b[4][13]), e) + expected_merge = templated_setindex!!(expected_merge, 2, @varname(e.b[4][14]), e) + @test @inferred(merge(vnt1, vnt2)) == expected_merge + + struct A + v # Just need something to disambiguate between two A's. + end + Base.size(::A) = (2,) + + @testset "with sized things: same indices" begin + # When things are stored as ArrayLikeBlocks, need to be careful that merging + # works correctly. There have been weird bugs in the past around this where + # part of an ArrayLikeBlock was set, and then the merged result had ALBs from + # both sides of the merge. + for (vn, template) in + [(@varname(x[1:2]), zeros(2)), (@varname(x[1:5][1:2]), zeros(5))] + # Regardless of how we set it above, the VNT should always store it as + # x[1:2]. + normalized_vn = @varname(x[1:2]) + + # Without templating. + vnt1 = @inferred(setindex!!(VarNamedTuple(), A(1.0), vn)) + vnt2 = @inferred(setindex!!(VarNamedTuple(), A(2.0), vn)) + @test vnt1[vn] == A(1.0) + @test vnt2[vn] == A(2.0) + merged = @inferred(merge(vnt1, vnt2)) + expected_merge = setindex!!(VarNamedTuple(), A(2.0), vn) + @test merged == expected_merge + @test only(keys(merged)) == normalized_vn + @test merged[vn] == A(2.0) + @test merged[normalized_vn] == A(2.0) + + # With templating. + vnt1 = @inferred( + templated_setindex!!(VarNamedTuple(), A(1.0), vn, template) + ) + @test vnt1[vn] == A(1.0) + vnt2 = @inferred( + templated_setindex!!(VarNamedTuple(), A(2.0), vn, template) + ) + @test vnt2[vn] == A(2.0) + merged = @inferred(merge(vnt1, vnt2)) + expected_merge = templated_setindex!!(VarNamedTuple(), A(2.0), vn, template) + @test merged == expected_merge + @test only(keys(merged)) == normalized_vn + @test merged[normalized_vn] == A(2.0) + end + end + + @testset "with sized things: different but overlapping indices" begin + vn1 = @varname(x[1:2]) + vn2 = @varname(x[2:3]) + + # Without templating. + vnt1 = @inferred(setindex!!(VarNamedTuple(), A(1.0), vn1)) + @test vnt1[vn1] == A(1.0) + vnt2 = @inferred(setindex!!(VarNamedTuple(), A(2.0), vn2)) + @test vnt2[vn2] == A(2.0) + expected_merge = setindex!!(VarNamedTuple(), A(2.0), vn2) + @test @inferred(merge(vnt1, vnt2)) == expected_merge + @test merge(vnt1, vnt2)[vn2] == A(2.0) + + # With templating. + vnt1 = @inferred(templated_setindex!!(VarNamedTuple(), A(1.0), vn1, zeros(3))) + @test vnt1[vn1] == A(1.0) + vnt2 = @inferred(templated_setindex!!(VarNamedTuple(), A(2.0), vn2, zeros(3))) + @test vnt2[vn2] == A(2.0) + merged = @inferred(merge(vnt1, vnt2)) + expected_merge = templated_setindex!!(VarNamedTuple(), A(2.0), vn2, zeros(3)) + @test merged == expected_merge + @test only(keys(merged)) == vn2 + @test merge(vnt1, vnt2)[vn2] == A(2.0) + end + + # TODO(penelopeysm): This set of tests fails. The reason is because later on we + # have a VarName that looks like d[1, 1][14, 13], i.e., d must be a matrix of + # matrices. So we have to provide that as the structure for d (via `f`). However, + # when you try to set `d[1, 3:4]` to `["1", "1"]` BangBang.setindex!! errors. + # https://github.com/JuliaFolds2/BangBang.jl/issues/44 + # + # I honestly think this is so pathological that it's not worth worrying about too + # much for now in DPPL. If someone complains about it, we can probably just honestly + # say that it's an upstream issue. + # + # Of course, I would like to fix it in BangBang, but time and energy, blah blah. + #= + d = (; d=fill(randn(20, 20), 4, 4)) + f = (; a=[(; b=(; c=fill(d, 4, 2)))]) + vnt1 = VarNamedTuple() + vnt1 = templated_setindex!!(vnt1, ["1", "1"], @varname(f.a[1].b.c[2, 2].d[1, 3:4]), f) + vnt2 = templated_setindex!!(vnt2, ["2", "2"], @varname(f.a[1].b.c[2, 2].d[1, 3:4]), f) + expected_merge = setindex!!( + expected_merge, ["2", "2"], @varname(f.a[1].b.c[2, 2].d[1, 3:4]) + ) + vnt1 = setindex!!(vnt1, :1, @varname(f.a[1].b.c[3, 2].d[1, 1][14, 13])) + vnt2 = setindex!!(vnt2, :2, @varname(f.a[1].b.c[4, 2].d[1, 1][14, 13])) + expected_merge = setindex!!( + expected_merge, :1, @varname(f.a[1].b.c[3, 2].d[1, 1][14, 13]) + ) + expected_merge = setindex!!( + expected_merge, :2, @varname(f.a[1].b.c[4, 2].d[1, 1][14, 13]) + ) + @test merge(vnt1, vnt2) == expected_merge + test_invariants(vnt1) + test_invariants(vnt2) + =# + end + + @testset "merging PartialArrays errors with different axes" begin + @testset "different sizes" begin + vnt1 = templated_setindex!!(VarNamedTuple(), 1.0, @varname(x[1]), zeros(1)) + vnt2 = templated_setindex!!(VarNamedTuple(), 2.0, @varname(x[1]), zeros(2)) + @test_throws ArgumentError merge(vnt1, vnt2) + + vnt1 = templated_setindex!!(VarNamedTuple(), 1.0, @varname(x[1]), zeros(1)) + vnt2 = templated_setindex!!(VarNamedTuple(), 2.0, @varname(x[1]), zeros(1, 1)) + @test_throws ArgumentError merge(vnt1, vnt2) + end + + @testset "different types" begin + vnt1 = templated_setindex!!(VarNamedTuple(), 1.0, @varname(x[1]), zeros(1)) + vnt2 = templated_setindex!!( + VarNamedTuple(), 2.0, @varname(x[0]), OA.OffsetArray(zeros(1), 0:0) + ) + @test_throws ArgumentError merge(vnt1, vnt2) + end + end + + @testset "_setindex_optic!! with MustNotOverwrite" begin + function test_must_not_overwrite(vnt, value, vn, template) + # Avoid mutating + vnt = deepcopy(vnt) + # Check that calling `_setindex_optic!!` with `MustNotOverwrite` errors if the + # variable already exists, and that it works if it doesn't. + @test_throws MustNotOverwriteError templated_setindex_no_overwrite!!( + vnt, value, vn, template + ) + @inferred(templated_setindex!!(vnt, value, vn, template)) + return nothing + end + + @testset "same variable" begin + vnt = templated_setindex!!(VarNamedTuple(), 1.0, @varname(x), NoTemplate()) + test_must_not_overwrite(vnt, 2.0, @varname(x), NoTemplate()) + + vnt = templated_setindex!!(VarNamedTuple(), 1.0, @varname(x.a), NoTemplate()) + test_must_not_overwrite(vnt, 2.0, @varname(x.a), NoTemplate()) + + vnt = templated_setindex!!(VarNamedTuple(), 1.0, @varname(x[1]), zeros(2)) + test_must_not_overwrite(vnt, 2.0, @varname(x[1]), zeros(2)) + + # Different indices, but the same slot + vnt = templated_setindex!!(VarNamedTuple(), 1.0, @varname(x[1]), zeros(2, 2)) + test_must_not_overwrite(vnt, 2.0, @varname(x[1, 1]), zeros(2, 2)) + end + + @testset "setting a superset of an old variable" begin + vnt = templated_setindex!!(VarNamedTuple(), 1.0, @varname(x.a), NoTemplate()) + test_must_not_overwrite(vnt, 2.0, @varname(x), NoTemplate()) + + vnt = templated_setindex!!(VarNamedTuple(), 1.0, @varname(x[1]), zeros(2)) + test_must_not_overwrite(vnt, 2.0, @varname(x), zeros(2)) + + vnt = templated_setindex!!(VarNamedTuple(), 1.0, @varname(x[1]), zeros(2)) + test_must_not_overwrite(vnt, [2.0, 3.0], @varname(x[1:2]), zeros(2)) + + vnt = templated_setindex!!(VarNamedTuple(), 1.0, @varname(x[1]), zeros(2)) + test_must_not_overwrite(vnt, [2.0, 3.0], @varname(x[:]), zeros(2)) + end + + @testset "setting a subset of an old variable" begin + vnt = templated_setindex!!(VarNamedTuple(), [1.0, 2.0], @varname(x), zeros(2)) + test_must_not_overwrite(vnt, 2.0, @varname(x[1]), zeros(2)) + + vnt = templated_setindex!!(VarNamedTuple(), [1.0, 2.0], @varname(x), zeros(2)) + test_must_not_overwrite(vnt, 2.0, @varname(x[end]), zeros(2)) + + vnt = templated_setindex!!(VarNamedTuple(), [1.0, 2.0], @varname(x), zeros(2)) + test_must_not_overwrite(vnt, [1.0, 2.0], @varname(x[1:2]), zeros(2)) + + vnt = templated_setindex!!( + VarNamedTuple(), [1.0, 2.0], @varname(x[1, :]), zeros(2, 2) + ) + test_must_not_overwrite(vnt, [1.0, 2.0], @varname(x[:, 1]), zeros(2, 2)) + + vnt = templated_setindex!!( + VarNamedTuple(), [1.0, 2.0], @varname(x.a), (; a=zeros(2)) + ) + test_must_not_overwrite(vnt, [1.0, 2.0], @varname(x.a[1:2]), (; a=zeros(2))) + + vnt = templated_setindex!!( + VarNamedTuple(), [1.0, 2.0], @varname(x.a), (; a=zeros(2)) + ) + test_must_not_overwrite(vnt, 2.0, @varname(x.a[1]), (; a=zeros(2))) + + vnt = templated_setindex!!( + VarNamedTuple(), [1.0, 2.0], @varname(x.a), (; a=zeros(2)) + ) + test_must_not_overwrite(vnt, 2.0, @varname(x.a[end]), (; a=zeros(2))) + + vnt = templated_setindex!!( + VarNamedTuple(), (; a=1.0), @varname(x), NoTemplate() + ) + test_must_not_overwrite(vnt, 2.0, @varname(x.a), NoTemplate()) + end + + @testset "different sub-indices of the same slice" begin + # Setting different sub-indices of the same slice should NOT error. + # This is the bug from issue #1321. + x = zeros(2) + vnt = templated_setindex_no_overwrite!!( + VarNamedTuple(), 1.0, @varname(x[1:2][1]), x + ) + vnt = templated_setindex_no_overwrite!!(vnt, 2.0, @varname(x[1:2][2]), x) + @test vnt[@varname(x)] == [1.0, 2.0] + + # But setting the SAME sub-index twice should still error. + vnt2 = templated_setindex_no_overwrite!!( + VarNamedTuple(), 1.0, @varname(x[1:2][1]), x + ) + test_must_not_overwrite(vnt2, 2.0, @varname(x[1:2][1]), x) + + # Also test with a larger array and different slices. + y = zeros(4) + vnt3 = templated_setindex_no_overwrite!!( + VarNamedTuple(), 1.0, @varname(y[1:3][1]), y + ) + vnt3 = templated_setindex_no_overwrite!!(vnt3, 2.0, @varname(y[1:3][2]), y) + test_must_not_overwrite(vnt3, [3.0, 4.0], @varname(y[1:3][2:3]), y) + vnt3 = templated_setindex_no_overwrite!!(vnt3, 3.0, @varname(y[1:3][3]), y) + @test vnt3[@varname(y[1:3])] == [1.0, 2.0, 3.0] + # Setting an already-set sub-index, or slice, should error. + test_must_not_overwrite(vnt3, 4.0, @varname(y[1:3][2]), y) + test_must_not_overwrite(vnt3, [4.0, 5.0], @varname(y[2:3][1:2]), y) + + # Also try with indices of different slices + z = zeros(3) + vnt4 = templated_setindex_no_overwrite!!( + VarNamedTuple(), [1.0, 2.0], @varname(z[1:2]), z + ) + test_must_not_overwrite(vnt4, 3.0, @varname(z[2:3][1]), z) + # but you can set it if it's not overlapping + templated_setindex_no_overwrite!!(vnt4, 3.0, @varname(z[2:3][2]), z) + end + end + + @testset "subset" begin + vnt = VarNamedTuple() + vnt = setindex!!(vnt, 1.0, @varname(a)) + vnt = setindex!!(vnt, [1, 2, 3], @varname(b)) + vnt = setindex!!(vnt, [10], @varname(c.x.y)) + d = randn(3) + vnt = templated_setindex!!(vnt, :1, @varname(d[1]), d) + vnt = setindex!!(vnt, :2, @varname(d[2])) + vnt = setindex!!(vnt, :3, @varname(d[3])) + e = (; f=fill((; g=(; h=fill((; i=0.0), 2, 4, 1))), 3, 3)) + vnt = templated_setindex!!(vnt, 2.0, @varname(e.f[3, 3].g.h[2, 4, 1].i), e) + p = fill(zeros(4, 5, 14), 2, 2) + vnt = templated_setindex!!( + vnt, SizedThing((3, 1, 4)), @varname(p[2, 1][2:4, 5:5, 11:14]), p + ) + test_invariants(vnt) + + # TODO(mhauru) I'm a bit saddened by the lack of type stability for subset: It's + # return type always infers as VarNamedTuple. Improving this would require a + # different implementation of subset. + @test subset(vnt, VarName[]) == VarNamedTuple() + @test subset(vnt, (@varname(z),)) == VarNamedTuple() + @test subset(vnt, (@varname(d[4]),)) == VarNamedTuple() + @test subset(vnt, (@varname(d[1, 1]),)) == VarNamedTuple() + @test subset(vnt, [@varname(a)]) == VarNamedTuple(; a=1.0) + begin + # These are a bit annoying to write out as we need to build the expected output + # manually + expected_vnt = setindex!!(VarNamedTuple(), [1, 2, 3], @varname(b)) + expected_vnt = templated_setindex!!(expected_vnt, :1, @varname(d[1]), randn(3)) + @test subset(vnt, [@varname(b), @varname(d[1])]) == expected_vnt + end + begin + expected_vnt = templated_setindex!!( + VarNamedTuple(), :2, @varname(d[2]), randn(3) + ) + expected_vnt = setindex!!(expected_vnt, :3, @varname(d[3])) + @test subset(vnt, [@varname(d[2:3])]) == expected_vnt + end + # For this one, we can't test exact equality because in the former case we + # constructed a PartialArray by adding individual elements of d. In the latter + # case we constructed it directly from a vector. So the internal representation + # is different, even though in the first case, the PartialArray has all its mask + # elements as true (and is thus conceptually the same as a full array). + # We can however test that the underlying values are the same when we try to extract + # them. + subsetted_d = subset(vnt, [@varname(d)]) + expected_d = VarNamedTuple((@varname(d) => [:1, :2, :3],)) + for vn in [ + @varname(d), + @varname(d[1]), + @varname(d[2]), + @varname(d[3]), + @varname(d[:]), + @varname(d[1:2]), + @varname(d[2:3]) + ] + @test getindex(subsetted_d, vn) == getindex(expected_d, vn) + end + + @test subset(vnt, [@varname(c.x.y)]) == VarNamedTuple((@varname(c.x.y) => [10],)) + @test subset(vnt, [@varname(c)]) == VarNamedTuple((@varname(c.x.y) => [10],)) + begin + expected_vnt = templated_setindex!!( + VarNamedTuple(), 2.0, @varname(e.f[3, 3].g.h[2, 4, 1].i), e + ) + @test subset(vnt, [@varname(e.f[3, 3].g.h[2, 4, 1].i)]) == expected_vnt + end + begin + expected_vnt = templated_setindex!!( + VarNamedTuple(), + SizedThing((3, 1, 4)), + @varname(p[2, 1][2:4, 5:5, 11:14]), + p, + ) + @test subset(vnt, [@varname(p[2, 1][2:4, 5:5, 11:14])]) == expected_vnt + end + # Cutting the last range a bit short should mean that nothing is returned. + @test subset(vnt, [@varname(p[2, 1][2:4, 5:5, 11:13])]) == VarNamedTuple() + end + + @testset "keys and values" begin + vnt = VarNamedTuple() + @test @inferred(keys(vnt)) == VarName[] + @test @inferred(values(vnt)) == Any[] + + vnt = setindex!!(vnt, 1.0, @varname(a)) + # TODO(mhauru) that the below passes @inferred, but any of the later ones don't. + # We should improve type stability of keys(). + @test @inferred(keys(vnt)) == [@varname(a)] + @test @inferred(values(vnt)) == [1.0] + + vnt = setindex!!(vnt, [1, 2, 3], @varname(b)) + @test keys(vnt) == [@varname(a), @varname(b)] + @test values(vnt) == [1.0, [1, 2, 3]] + + vnt = setindex!!(vnt, 15, @varname(b[2])) + @test keys(vnt) == [@varname(a), @varname(b)] + @test values(vnt) == [1.0, [1, 15, 3]] + + vnt = setindex!!(vnt, [10], @varname(c.x.y)) + @test keys(vnt) == [@varname(a), @varname(b), @varname(c.x.y)] + @test values(vnt) == [1.0, [1, 15, 3], [10]] + + vnt = setindex!!(vnt, -1.0, @varname(d[4])) + @test keys(vnt) == [@varname(a), @varname(b), @varname(c.x.y), @varname(d[4])] + @test values(vnt) == [1.0, [1, 15, 3], [10], -1.0] + + vnt = setindex!!(vnt, 2.0, @varname(e.f[3, 3].g.h[2, 4, 1].i)) + @test keys(vnt) == [ + @varname(a), + @varname(b), + @varname(c.x.y), + @varname(d[4]), + @varname(e.f[3, 3].g.h[2, 4, 1].i), + ] + @test values(vnt) == [1.0, [1, 15, 3], [10], -1.0, 2.0] + + vnt = setindex!!(vnt, fill(1.0, 4), @varname(j[1:4])) + @test keys(vnt) == [ + @varname(a), + @varname(b), + @varname(c.x.y), + @varname(d[4]), + @varname(e.f[3, 3].g.h[2, 4, 1].i), + @varname(j[1]), + @varname(j[2]), + @varname(j[3]), + @varname(j[4]), + ] + @test values(vnt) == [1.0, [1, 15, 3], [10], -1.0, 2.0, fill(1.0, 4)...] + + vnt = setindex!!(vnt, "a", @varname(j[6])) + @test keys(vnt) == [ + @varname(a), + @varname(b), + @varname(c.x.y), + @varname(d[4]), + @varname(e.f[3, 3].g.h[2, 4, 1].i), + @varname(j[1]), + @varname(j[2]), + @varname(j[3]), + @varname(j[4]), + @varname(j[6]), + ] + @test values(vnt) == [1.0, [1, 15, 3], [10], -1.0, 2.0, fill(1.0, 4)..., "a"] + + vnt = setindex!!(vnt, 1.0, @varname(n[2].a)) + @test keys(vnt) == [ + @varname(a), + @varname(b), + @varname(c.x.y), + @varname(d[4]), + @varname(e.f[3, 3].g.h[2, 4, 1].i), + @varname(j[1]), + @varname(j[2]), + @varname(j[3]), + @varname(j[4]), + @varname(j[6]), + @varname(n[2].a), + ] + @test values(vnt) == [1.0, [1, 15, 3], [10], -1.0, 2.0, fill(1.0, 4)..., "a", 1.0] + + vnt = setindex!!(vnt, SizedThing((3, 1, 4)), @varname(o[2:4, 5:5, 11:14])) + @test keys(vnt) == [ + @varname(a), + @varname(b), + @varname(c.x.y), + @varname(d[4]), + @varname(e.f[3, 3].g.h[2, 4, 1].i), + @varname(j[1]), + @varname(j[2]), + @varname(j[3]), + @varname(j[4]), + @varname(j[6]), + @varname(n[2].a), + @varname(o[2:4, 5:5, 11:14]), + ] + @test values(vnt) == [ + 1.0, + [1, 15, 3], + [10], + -1.0, + 2.0, + fill(1.0, 4)..., + "a", + 1.0, + SizedThing((3, 1, 4)), + ] + + vnt = setindex!!(vnt, SizedThing((3, 1, 4)), @varname(p[2, 1][2:4, 5:5, 11:14])) + @test keys(vnt) == [ + @varname(a), + @varname(b), + @varname(c.x.y), + @varname(d[4]), + @varname(e.f[3, 3].g.h[2, 4, 1].i), + @varname(j[1]), + @varname(j[2]), + @varname(j[3]), + @varname(j[4]), + @varname(j[6]), + @varname(n[2].a), + @varname(o[2:4, 5:5, 11:14]), + @varname(p[2, 1][2:4, 5:5, 11:14]), + ] + @test values(vnt) == [ + 1.0, + [1, 15, 3], + [10], + -1.0, + 2.0, + fill(1.0, 4)..., + "a", + 1.0, + SizedThing((3, 1, 4)), + SizedThing((3, 1, 4)), + ] + test_invariants(vnt) + end + + @testset "length" begin + # Type inference for length fails in some cases on Julia versions < 1.11 + inference_broken = VERSION < v"1.11" + + vnt = VarNamedTuple() + @test @inferred(length(vnt)) == 0 + + vnt = setindex!!(vnt, 1.0, @varname(a)) + @test @inferred(length(vnt)) == 1 + + vnt = setindex!!(vnt, [1, 2, 3], @varname(b)) + @test @inferred(length(vnt)) == 2 + + vnt = setindex!!(vnt, 15, @varname(b[2])) + @test @inferred(length(vnt)) == 2 + + vnt = setindex!!(vnt, [10, 11], @varname(c.x.y)) + @test @inferred(length(vnt)) == 3 + + vnt = setindex!!(vnt, -1.0, @varname(d[4])) + @test @inferred(length(vnt)) == 4 broken = inference_broken + + vnt = setindex!!(vnt, ["a", "b"], @varname(d[1:2])) + @test @inferred(length(vnt)) == 6 broken = inference_broken + + vnt = setindex!!(vnt, 2.0, @varname(e.f[3].g.h[2].i)) + vnt = setindex!!(vnt, 3.0, @varname(e.f[3].g.h[2].j)) + @test @inferred(length(vnt)) == 8 broken = inference_broken + + vnt = setindex!!(vnt, SizedThing((3, 2)), @varname(x[1, 2:4, 2, 1:2, 3])) + @test @inferred(length(vnt)) == 14 broken = inference_broken + + vnt = setindex!!(vnt, SizedThing((3, 2)), @varname(x[1, 4:6, 2, 1:2, 3])) + @test @inferred(length(vnt)) == 14 broken = inference_broken + + vnt = setindex!!(vnt, [:a, :b], @varname(y[4][3][2][1:2])) + @test @inferred(length(vnt)) == 16 broken = inference_broken + test_invariants(vnt) + end + + @testset "empty" begin + # test_invariants already checks that many different kinds of VarNamedTuples can be + # emptied with empty and empty!!. What remains to check here is that + # 1) isempty gives the expected results: + vnt = VarNamedTuple() + @test @inferred(isempty(vnt)) == true + vnt = setindex!!(vnt, 1.0, @varname(a)) + @test @inferred(isempty(vnt)) == false + test_invariants(vnt) + + # Setting an empty value inside a VNT doesn't mean that the VNT is empty + vnt = VarNamedTuple() + vnt = templated_setindex!!(vnt, [], @varname(a[1]), zeros(1)) + @test @inferred(isempty(vnt)) == false + test_invariants(vnt) + + # 2) empty!! keeps PartialArrays in place: + vnt = VarNamedTuple() + vnt = @inferred(templated_setindex!!(vnt, [1, 2, 3], @varname(a[1:3]), zeros(4))) + vnt = @inferred(empty!!(vnt)) + @test !haskey(vnt, @varname(a[1])) + @test !haskey(vnt, @varname(a[1:3])) + @test !haskey(vnt, @varname(a)) + # It's an empty PartialArray now, so attempting to extract anything will fail + @test_throws BoundsError getindex(vnt, @varname(a[1])) + @test_throws BoundsError getindex(vnt, @varname(a[1:3])) + @test_throws ArgumentError getindex(vnt, @varname(a)) + # but checking it explicitly, it should still be there + @test haskey(vnt.data, :a) + @test vnt.data.a isa PartialArray + @test !any(vnt.data.a.mask) + # We can set data back into it and it will reuse the same template + vnt = @inferred(setindex!!(vnt, [1, 2, 3], @varname(a[2:4]))) + @test @inferred(getindex(vnt, @varname(a[2:4]))) == [1, 2, 3] + @test haskey(vnt, @varname(a[2:4])) + @test !haskey(vnt, @varname(a[1])) + test_invariants(vnt) + end + + @testset "densification" begin + vnt = VarNamedTuple() + vnt = @inferred(setindex!!(vnt, 1.0, @varname(a.b[1].c[1, 1]))) + @test @inferred(getindex(vnt, @varname(a.b[1].c))) == fill(1.0, (1, 1)) + test_invariants(vnt) + + vnt = VarNamedTuple() + vnt = @inferred(setindex!!(vnt, 1.0, @varname(a.b[1].c[1, 1]))) + vnt = @inferred(setindex!!(vnt, 1.0, @varname(a.b[1].c[1, 2]))) + @test @inferred(getindex(vnt, @varname(a.b[1].c))) == fill(1.0, (1, 2)) + test_invariants(vnt) + + vnt = VarNamedTuple() + vnt = @inferred(setindex!!(vnt, 1.0, @varname(a.b[1].c[1, 1]))) + vnt = @inferred(setindex!!(vnt, 1.0, @varname(a.b[1].c[2, 1]))) + @test @inferred(getindex(vnt, @varname(a.b[1].c))) == fill(1.0, (2, 1)) + test_invariants(vnt) + + vnt = VarNamedTuple() + vnt = @inferred(setindex!!(vnt, 1.0, @varname(a.b[1].c[1, 1]))) + vnt = @inferred(setindex!!(vnt, 1.0, @varname(a.b[1].c[1, 2]))) + vnt = @inferred(setindex!!(vnt, 1.0, @varname(a.b[1].c[2, 1]))) + @test_throws ArgumentError @inferred(getindex(vnt, @varname(a.b[1].c))) + vnt = @inferred(setindex!!(vnt, 1.0, @varname(a.b[1].c[2, 2]))) + @test @inferred(getindex(vnt, @varname(a.b[1].c))) == fill(1.0, (2, 2)) + vnt = @inferred(setindex!!(vnt, 1.0, @varname(a.b[1].c[3, 3]))) + @test_throws ArgumentError @inferred(getindex(vnt, @varname(a.b[1].c))) + test_invariants(vnt) + + vnt = VarNamedTuple() + vnt = @inferred(setindex!!(vnt, SizedThing((2,)), @varname(x[1:2]))) + @test_throws ArgumentError @inferred(getindex(vnt, @varname(x))) + test_invariants(vnt) + end + + @testset "block variables" begin + # Tests for setting and getting block variables, i.e. variables that have a non-zero + # size in a PartialArray, but are not Arrays themselves. + expected_err = ArgumentError(""" + A non-Array value set with a range of indices must be retrieved with the same + range of indices. + """) + vnt = VarNamedTuple() + vnt = @inferred(setindex!!(vnt, SizedThing((3,)), @varname(x[2:4]))) + test_invariants(vnt) + @test haskey(vnt, @varname(x[2:4])) + @test @inferred(getindex(vnt, @varname(x[2:4]))) == SizedThing((3,)) + @test !haskey(vnt, @varname(x[2:3])) + @test_throws expected_err getindex(vnt, @varname(x[2:3])) + @test !haskey(vnt, @varname(x[3])) + @test_throws expected_err getindex(vnt, @varname(x[3])) + @test !haskey(vnt, @varname(x[1])) + @test !haskey(vnt, @varname(x[5])) + vnt = setindex!!(vnt, 1.0, @varname(x[1])) + vnt = setindex!!(vnt, 1.0, @varname(x[5])) + test_invariants(vnt) + @test haskey(vnt, @varname(x[1])) + @test haskey(vnt, @varname(x[5])) + @test_throws expected_err getindex(vnt, @varname(x[1:4])) + @test_throws expected_err getindex(vnt, @varname(x[2:5])) + + # Setting any of these indices should remove the block variable x[2:4]. + @testset "index = $index" for index in (2, 3, 4, 2:3, 3:5) + # Test setting different types of values. + vals = if index isa Int + (2.0,) + else + (fill(2.0, length(index)), SizedThing((length(index),))) + end + @testset "val = $val" for val in vals + vn = @varname(x[index]) + vnt2 = copy(vnt) + vnt2 = setindex!!(vnt2, val, vn) + test_invariants(vnt) + @test !haskey(vnt2, @varname(x[2:4])) + @test_throws BoundsError getindex(vnt2, @varname(x[2:4])) + other_index = index in (2, 2:3) ? 4 : 2 + @test !haskey(vnt2, @varname(x[other_index])) + @test_throws BoundsError getindex(vnt2, @varname(x[other_index])) + @test haskey(vnt2, vn) + @test getindex(vnt2, vn) == val + @test haskey(vnt2, @varname(x[1])) + @test_throws BoundsError getindex(vnt2, @varname(x[1:4])) + end + end + + # Extra checks, mostly for type stability and to confirm that multidimensional + # blocks work too. + val = SizedThing((2, 2)) + vnt = VarNamedTuple() + vnt = @inferred(setindex!!(vnt, val, @varname(y.z[1:2, 1:2]))) + test_invariants(vnt) + @test haskey(vnt, @varname(y.z[1:2, 1:2])) + @test @inferred(getindex(vnt, @varname(y.z[1:2, 1:2]))) == val + @test !haskey(vnt, @varname(y.z[1, 1])) + @test_throws expected_err getindex(vnt, @varname(y.z[1, 1])) + + vnt = @inferred(setindex!!(vnt, val, @varname(y.z[2:3, 2:3]))) + test_invariants(vnt) + @test haskey(vnt, @varname(y.z[2:3, 2:3])) + @test @inferred(getindex(vnt, @varname(y.z[2:3, 2:3]))) == val + @test !haskey(vnt, @varname(y.z[1:2, 1:2])) + @test_throws BoundsError getindex(vnt, @varname(y.z[1:2, 1:2])) + + vnt = @inferred(setindex!!(vnt, val, @varname(y.z[4:5, 2:3]))) + test_invariants(vnt) + @test haskey(vnt, @varname(y.z[2:3, 2:3])) + @test @inferred(getindex(vnt, @varname(y.z[2:3, 2:3]))) == val + @test haskey(vnt, @varname(y.z[4:5, 2:3])) + @test @inferred(getindex(vnt, @varname(y.z[4:5, 2:3]))) == val + + # A lot like above, but with extra indices that are not ranges. + val = SizedThing((2, 2)) + vnt = VarNamedTuple() + vnt = @inferred(setindex!!(vnt, val, @varname(y.z[2, 1:2, 3, 1:2, 4]))) + test_invariants(vnt) + @test haskey(vnt, @varname(y.z[2, 1:2, 3, 1:2, 4])) + @test @inferred(getindex(vnt, @varname(y.z[2, 1:2, 3, 1:2, 4]))) == val + @test !haskey(vnt, @varname(y.z[2, 1, 3, 1, 4])) + @test_throws expected_err getindex(vnt, @varname(y.z[2, 1, 3, 1, 4])) + + vnt = @inferred(setindex!!(vnt, val, @varname(y.z[2, 2:3, 3, 2:3, 4]))) + test_invariants(vnt) + @test haskey(vnt, @varname(y.z[2, 2:3, 3, 2:3, 4])) + @test @inferred(getindex(vnt, @varname(y.z[2, 2:3, 3, 2:3, 4]))) == val + @test !haskey(vnt, @varname(y.z[2, 1:2, 3, 1:2, 4])) + @test_throws BoundsError getindex(vnt, @varname(y.z[2, 1:2, 3, 1:2, 4])) + + vnt = @inferred(setindex!!(vnt, val, @varname(y.z[3, 2:3, 3, 2:3, 4]))) + test_invariants(vnt) + @test haskey(vnt, @varname(y.z[2, 2:3, 3, 2:3, 4])) + @test @inferred(getindex(vnt, @varname(y.z[2, 2:3, 3, 2:3, 4]))) == val + @test haskey(vnt, @varname(y.z[3, 2:3, 3, 2:3, 4])) + # Type inference fails on this one for Julia versions < 1.11 + @test @inferred(getindex(vnt, @varname(y.z[3, 2:3, 3, 2:3, 4]))) == val + end + + @testset "map and friends" begin + vnt = VarNamedTuple() + vnt = @inferred(setindex!!(vnt, 1, @varname(a))) + vnt = @inferred(setindex!!(vnt, [2, 2], @varname(b[1:2]))) + vnt = @inferred(setindex!!(vnt, [3.0], @varname(c.d))) + vnt = @inferred(setindex!!(vnt, "a", @varname(e.f[3].g.h[2].i))) + # The below can't be type stable because the element type of `h` depends on whether + # we are setting `h[2].j` (which overwrites the earlier `h[2]`) or some other + # `h[index].j` (which would leave both `h[2].i` and `h[index].j` in the same array). + vnt = setindex!!(vnt, 5.0, @varname(e.f[3].g.h[2].j)) + vnt = @inferred( + setindex!!(vnt, SizedThing((2, 2)), @varname(y.z[3, 2:3, 3, 2:3, 4])) + ) + vnt = @inferred( + templated_setindex!!(vnt, SizedThing((2,)), @varname(v[:]), randn(2)) + ) + vnt = @inferred(setindex!!(vnt, "", @varname(w[4][3][2, 1]))) + test_invariants(vnt) + + struct AnotherSizedThing{T<:Tuple} + size::T + end + Base.size(st::AnotherSizedThing) = st.size + + call_counter = 0 + function f_val(val) + call_counter += 1 + if val isa Int + return val + 10 + elseif val isa AbstractVector{Int} + return val .+ 10 + elseif val isa Float64 + return val + 1.0 + elseif val isa AbstractVector{Float64} + return val .- 1.0 + elseif val isa String + return string(val, "b") + elseif val isa SizedThing + return AnotherSizedThing(size(val)) + else + error("Unexpected value type $(typeof(val))") + end + end + + f_pair(pair) = f_val(pair.second) + + val_reduction = mapreduce(pair -> pair.second, vcat, vnt; init=Any[]) + @test val_reduction == vcat( + Any[], 1, [2, 2], [3.0], "a", 5.0, SizedThing((2, 2)), SizedThing((2,)), "" + ) + key_reduction = mapreduce(pair -> pair.first, vcat, vnt; init=Any[]) + @test key_reduction == vcat( + @varname(a), + @varname(b[1]), + @varname(b[2]), + @varname(c.d), + @varname(e.f[3].g.h[2].i), + @varname(e.f[3].g.h[2].j), + @varname(y.z[3, 2:3, 3, 2:3, 4]), + @varname(v[:]), + @varname(w[4][3][2, 1]), + ) + + call_counter = 0 + reduction = mapreduce(f_pair, vcat, vnt; init=Any[]) + @test reduction == vcat( + Any[], + 11, + [12, 12], + [2.0], + "ab", + 6.0, + AnotherSizedThing((2, 2)), + AnotherSizedThing((2,)), + "b", + ) + # Check that f_pair gets called exactly once per element. + @test call_counter == length(keys(vnt)) + + # TODO(mhauru) This should hopefully be type stable, but fails to be so because of + # some complex VarNames being too much for constant propagation. See comment in + # src/varnamedtuple.jl for more. + call_counter = 0 + vnt_mapped = map_pairs!!(f_pair, copy(vnt)) + # Check that f_pair gets called exactly once per element. + @test call_counter == length(keys(vnt)) + @test vnt_mapped == map_values!!(f_val, copy(vnt)) + test_invariants(vnt_mapped) + @test @inferred(getindex(vnt_mapped, @varname(a))) == 11 + @test @inferred(getindex(vnt_mapped, @varname(b[1:2]))) == [12, 12] + @test @inferred(getindex(vnt_mapped, @varname(c.d))) == [2.0] + @test @inferred(getindex(vnt_mapped, @varname(e.f[3].g.h[2].i))) == "ab" + @test @inferred(getindex(vnt_mapped, @varname(e.f[3].g.h[2].j))) == 6.0 + @test @inferred(getindex(vnt_mapped, @varname(y.z[3, 2:3, 3, 2:3, 4]))) == + AnotherSizedThing((2, 2)) + @test @inferred(getindex(vnt_mapped, @varname(v[:]))) == AnotherSizedThing((2,)) + @test @inferred(getindex(vnt_mapped, @varname(w[4][3][2, 1]))) == "b" + + call_counter = 0 + vnt_applied = copy(vnt) + vnt_applied = @inferred(apply!!(f_val, vnt_applied, @varname(a))) + @test call_counter == 1 + test_invariants(vnt_applied) + @test @inferred(getindex(vnt_applied, @varname(a))) == 11 + @test @inferred(getindex(vnt_applied, @varname(b[1:2]))) == [2, 2] + + vnt_applied = @inferred(apply!!(f_val, vnt_applied, @varname(b[1:2]))) + # Unlike map_pairs!!, apply!! operates on the whole value at once, rather than + # element-wise, so this is only one more call. + @test call_counter == 2 + test_invariants(vnt_applied) + @test @inferred(getindex(vnt_applied, @varname(a))) == 11 + @test @inferred(getindex(vnt_applied, @varname(b[1:2]))) == [12, 12] + + vnt_applied = @inferred(apply!!(f_val, vnt_applied, @varname(c.d))) + @test call_counter == 3 + test_invariants(vnt_applied) + @test @inferred(getindex(vnt_applied, @varname(c.d))) == [2.0] + + vnt_applied = begin + # The @inferred fails on Julia 1.10. + @static if VERSION < v"1.11" + apply!!(f_val, vnt_applied, @varname(e.f[3].g.h[2].i)) + else + @inferred(apply!!(f_val, vnt_applied, @varname(e.f[3].g.h[2].i))) + end + end + @test call_counter == 4 + test_invariants(vnt_applied) + @test @inferred(getindex(vnt_applied, @varname(e.f[3].g.h[2].i))) == "ab" + @test @inferred(getindex(vnt_applied, @varname(e.f[3].g.h[2].j))) == 5.0 + + vnt_applied = begin + # The @inferred fails on Julia 1.10. + @static if VERSION < v"1.11" + apply!!(f_val, vnt_applied, @varname(e.f[3].g.h[2].j)) + else + @inferred(apply!!(f_val, vnt_applied, @varname(e.f[3].g.h[2].j))) + end + end + @test call_counter == 5 + test_invariants(vnt_applied) + @test @inferred(getindex(vnt_applied, @varname(e.f[3].g.h[2].j))) == 6.0 + + # This can't be type stable because y.z might have many elements set, and we can't + # know at compile time that this sets the only one, thus allowing the element type + # to be AnotherSizedThing. + vnt_applied = apply!!(f_val, vnt_applied, @varname(y.z[3, 2:3, 3, 2:3, 4])) + @test call_counter == 6 + test_invariants(vnt_applied) + @test @inferred(getindex(vnt_applied, @varname(y.z[3, 2:3, 3, 2:3, 4]))) == + AnotherSizedThing((2, 2)) + + vnt_applied = apply!!(f_val, vnt_applied, @varname(v[:])) + @test call_counter == 7 + test_invariants(vnt_applied) + @test @inferred(getindex(vnt_applied, @varname(v[:]))) == AnotherSizedThing((2,)) + + vnt_applied = @inferred(apply!!(f_val, vnt_applied, @varname(w[4][3][2, 1]))) + @test call_counter == 8 + test_invariants(vnt_applied) + @test @inferred(getindex(vnt_applied, @varname(w[4][3][2, 1]))) == "b" + + # map a function that maps every key => value pair to key => key. + # For this, use a simpler VarNamedTuple, because block variables don't work with + # this mapping function. It also allows us to check type stability. + vnt = VarNamedTuple() + vnt = @inferred(setindex!!(vnt, 1, @varname(a))) + vnt = @inferred(setindex!!(vnt, 2, @varname(b[2]))) + vnt = @inferred(setindex!!(vnt, [3.0], @varname(c.d))) + vnt = @inferred(setindex!!(vnt, :oi, @varname(y.z[3, 2, 3, 2, 4]))) + vnt = @inferred(setindex!!(vnt, "", @varname(w[4][2, 1]))) + + get_key(pair) = pair.first + vnt_key_mapped = @inferred(map_pairs!!(get_key, copy(vnt))) + vnt_key_mapped_expected = VarNamedTuple() + for k in keys(vnt) + vnt_key_mapped_expected = setindex!!(vnt_key_mapped_expected, k, k) + end + @test vnt_key_mapped == vnt_key_mapped_expected + end + + @testset "Conversion to NamedTuple" begin + @testset "type stability when it works" begin + vnt = AbstractPPL.VarNamedTuple() + @test @inferred(NamedTuple(vnt)) == NamedTuple() + vnt = setindex!!(vnt, 1.0, @varname(a)) + @test @inferred(NamedTuple(vnt)) == (a=1.0,) + vnt = setindex!!(vnt, [1, 2, 3], @varname(b)) + @test @inferred(NamedTuple(vnt)) == (a=1.0, b=[1, 2, 3]) + vnt = setindex!!(vnt, "wut", @varname(b[2])) + @test @inferred(NamedTuple(vnt)) == (a=1.0, b=[1, "wut", 3]) + end + + @testset "invalid VarNames" begin + vnt = AbstractPPL.VarNamedTuple() + vnt = setindex!!(vnt, 1.0, @varname(a[1])) + @test_throws ArgumentError NamedTuple(vnt) + + vnt = AbstractPPL.VarNamedTuple() + vnt = setindex!!(vnt, 1.0, @varname(a.b)) + @test_throws ArgumentError NamedTuple(vnt) + end + end + + @testset "@vnt macro" begin + @testset "no templates" begin + vnt = @vnt begin + a := 1.0 + b := [1, 15, 3] + c.x.y := [10] + end + expected_vnt = VarNamedTuple() + expected_vnt = setindex!!(expected_vnt, 1.0, @varname(a)) + expected_vnt = setindex!!(expected_vnt, [1, 15, 3], @varname(b)) + expected_vnt = setindex!!(expected_vnt, [10], @varname(c.x.y)) + @test vnt == expected_vnt + end + + @testset "rhs values outside of vnt block" begin + a = 1.0 + b = [1, 15, 3] + cxy = [10] + vnt = @vnt begin + a := a + b := b + c.x.y := cxy + end + expected_vnt = VarNamedTuple() + expected_vnt = setindex!!(expected_vnt, 1.0, @varname(a)) + expected_vnt = setindex!!(expected_vnt, [1, 15, 3], @varname(b)) + expected_vnt = setindex!!(expected_vnt, [10], @varname(c.x.y)) + @test vnt == expected_vnt + end + + @testset "templates as top-level symbols" begin + a = zeros(2, 2) + vnt = @vnt begin + @template a + a[1] := 1.0 + end + expected_vnt = VarNamedTuple() + expected_vnt = templated_setindex!!(expected_vnt, 1.0, @varname(a[1]), a) + @test vnt == expected_vnt + end + + @testset "scoping of templates" begin + a = zeros(2, 2) # This should be ignored + vnt = begin + a = zeros(3) + @vnt begin + @template a + a[1] := 1.0 + end + end + expected_vnt = VarNamedTuple() + expected_vnt = templated_setindex!!(expected_vnt, 1.0, @varname(a[1]), zeros(3)) + @test vnt == expected_vnt + end + + @testset "templates with assignment" begin + a = zeros(3) + vnt = @vnt begin + @template b = a + b[1] := 2.0 + end + expected_vnt = VarNamedTuple() + expected_vnt = templated_setindex!!(expected_vnt, 2.0, @varname(b[1]), a) + @test vnt == expected_vnt + end + + @testset "templates with expr" begin + evalcount = Ref(0) + function make_template() + evalcount[] += 1 + return zeros(4) + end + vnt = @vnt begin + @template c = make_template() + c[1] := 1.0 + c[2] := 2.0 + end + # Check that the template expression is only evaluated once. + @test evalcount[] == 1 + + expected_vnt = VarNamedTuple() + expected_vnt = templated_setindex!!(expected_vnt, 1.0, @varname(c[1]), zeros(4)) + expected_vnt = templated_setindex!!(expected_vnt, 2.0, @varname(c[2]), zeros(4)) + @test vnt == expected_vnt + end + end + + @testset "densify!!" begin + vnt = @vnt begin + @template x = zeros(2) + x[1] := 1.0 + x[2] := 2.0 + end + @test densify!!(vnt) == VarNamedTuple(; x=[1.0, 2.0]) + + # Nested + vnt = @vnt begin + @template x = fill(zeros(2), 2) + x[1][1] := 1.0 + x[2][1] := 2.0 + x[1][2] := 3.0 + x[2][2] := 4.0 + end + @test densify!!(vnt) == VarNamedTuple(; x=[[1.0, 3.0], [2.0, 4.0]]) + + # Check that it's type stable when there are no PAs. + vnt = @vnt begin + x := [1.0, 2.0] + y := 3.0 + z := "a" + end + @test @inferred(densify!!(vnt)) == vnt + + # Check that it doesn't densify GrowableArrays (and that failing to do so is type + # stable). + vnt = @vnt begin + x[1] := 1.0 + x[2] := 2.0 + end + @test @inferred(densify!!(vnt)) == vnt + + # Check that it doesn't densify PAs that have VNTs. Note that this is not type + # stable because it needs to recurse into the VNT to determine whether the VNT + # itself has PAs that need to be densified. + vnt = @vnt begin + @template x = zeros(2) + x[1].a := 1.0 + x[2].b := 2.0 + end + @test densify!!(vnt) == vnt + + # Check that densify!! recurses into VNTs inside PAs (VNT -> PA -> VNT -> PA) + vnt = @vnt begin + @template x = fill((; y=zeros(2)), 1) + x[1].y[1] := 1.0 + x[1].y[2] := 2.0 + end + result = densify!!(vnt) + # Outer PA (holding VNTs) should remain a PartialArray + @test result.data.x isa PartialArray + # Inner PA (holding Floats, fully filled) should be densified to a plain Vector + @test result.data.x[1].data.y == [1.0, 2.0] + @test !(result.data.x[1].data.y isa PartialArray) + + # Check that it doesn't densify PAs that have ALBs. + # Note: this isn't type stable because it must recurse into each element of the PA + # to check for densification opportunities. + # This might seem silly, because neither Float64 and ALBs can be densified. The + # problem here is that `eltype(pa.data)` (a mixture of Float64 and ALBs) is not + # Union{Float64,ALB}, but rather just Any. Because of that, the compiler can't + # statically determine that there aren't any e.g. VNTs that might themselves contain + # PAs that need to be densified. + vnt = @vnt begin + @template x = zeros(3) + x[1] := 1.0 + x[2:3] := SizedThing((2,)) + end + @test densify!!(vnt) == vnt + + # Check with ComponentArrays + x = CA.ComponentArray(; a=0.0, b=0.0) + vnt = @vnt begin + @template x + x.a := 1.0 + x.b := 2.0 + end + @test densify!!(vnt) == VarNamedTuple(; x=CA.ComponentArray(; a=1.0, b=2.0)) + end + + @testset "skeleton" begin + function test_skeleton(orig_vnt, expected_skeleton) + @test (@inferred skeleton(orig_vnt)) == expected_skeleton + skel = skeleton(orig_vnt) + # Check roundtrip reconstruction + new_vnt = VarNamedTuple() + for (vn, val) in pairs(orig_vnt) + top_sym = AbstractPPL.getsym(vn) + template = get(skel.data, top_sym, AbstractPPL.NoTemplate()) + new_vnt = AbstractPPL.templated_setindex!!(new_vnt, val, vn, template) + end + @test new_vnt == orig_vnt + end + + # Empty + v0 = VarNamedTuple() + v0s = VarNamedTuple() + test_skeleton(v0, v0s) + + # VNT + v1 = @vnt begin + y := 2.0 + end + v1s = VarNamedTuple() + test_skeleton(v1, v1s) + + # VNT -> VNT -> VNT -> ... but only VNTs + v2 = @vnt begin + y := 2.0 + z.a := 3.0 + a.b.c.d.e.f := [1.0, 2.0, 3.0] + c := "string" + end + v2s = VarNamedTuple() + test_skeleton(v2, v2s) + + # VNT -> PA[Float] + v3 = @vnt begin + @template x = zeros(2) + x[1] := 1.0 + end + v3s = VarNamedTuple(; x=fill(nothing, 2)) + test_skeleton(v3, v3s) + + # VNT -> VNT -> PA + v4 = @vnt begin + @template x = (; y=zeros(2)) + x.y[2] := 2.0 + end + v4s = VarNamedTuple(; x=VarNamedTuple(; y=fill(nothing, 2))) + test_skeleton(v4, v4s) + + # VNT -> PA -> VNT + v5 = @vnt begin + @template x = zeros(2) + x[1].y := 2.0 + end + v5s = VarNamedTuple(; x=fill(nothing, 2)) + test_skeleton(v5, v5s) + + # VNT -> PA -> VNT -> VNT + v6 = @vnt begin + @template x = fill((; y=3.0), 2) + x[1].y.z := "wut" + end + v6s = VarNamedTuple(; x=fill(nothing, 2)) + test_skeleton(v6, v6s) + + # VNT -> PA -> PA + v7 = @vnt begin + @template x = fill(zeros(2), 3) + x[3][2] := 2.0 + end + v7s = VarNamedTuple(; x=fill(fill(nothing, 2), 3)) + test_skeleton(v7, v7s) + + # VNT -> PA -> VNT -> PA + v8 = @vnt begin + @template x = fill((; y=zeros(2)), 3) + x[3].y[2] := 2.0 + end + v8s = VarNamedTuple(; x=fill(VarNamedTuple(; y=fill(nothing, 2)), 3)) + test_skeleton(v8, v8s) + + # VNT -> PA[Any] + v9 = @vnt begin + @template x = zeros(3) + x[1] := 1.0 + x[2] := "wut" + end + v9s = VarNamedTuple(; x=fill(nothing, 3)) + test_skeleton(v9, v9s) + + # VNT -> PA with some entries needing to recurse + v10 = @vnt begin + @template x = fill((; y=zeros(2)), 3) + x[3].y[2] := 2.0 + x[1] := 1.0 + end + v10s = VarNamedTuple(; x=[nothing, nothing, VarNamedTuple(; y=fill(nothing, 2))]) + test_skeleton(v10, v10s) + + # VNT -> PA with different types of arrays + v11 = @vnt begin + @template x = OA.OffsetArray(zeros(3), -4:-2) + x[-2] := 2.0 + end + v11s = VarNamedTuple(; x=OA.OffsetArray(fill(nothing, 3), -4:-2)) + test_skeleton(v11, v11s) + + v12 = @vnt begin + @template x = DD.DimArray(zeros(2, 3), (:a, :b)) + x[1, 2] := 2.0 + end + v12s = VarNamedTuple(; x=DD.DimArray(fill(nothing, 2, 3), (:a, :b))) + test_skeleton(v12, v12s) + + v13 = @vnt begin + @template x = CA.ComponentArray(; a=0.0, b=0.0) + x.a := 1.0 + end + v13s = VarNamedTuple(; x=CA.ComponentArray(; a=nothing, b=nothing)) + test_skeleton(v13, v13s) + end +end + +end From 4a71addfe5e3315c25343f944da4dff81cadf0d3 Mon Sep 17 00:00:00 2001 From: Xianda Sun Date: Mon, 13 Jul 2026 08:29:39 +0100 Subject: [PATCH 2/3] Release the VarNamedTuple migration as 0.16.0; formatting fixes --- HISTORY.md | 12 +++++++----- Project.toml | 2 +- src/varnamedtuple/getset.jl | 2 +- src/varnamedtuple/vnt.jl | 19 +++++++++++-------- test/varname/hasvalue.jl | 22 +++++++++++----------- 5 files changed, 31 insertions(+), 26 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index 80f312a3..6933204d 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,8 +1,4 @@ -## 0.15.4 - -`value_and_gradient!!`, `value_and_jacobian!!`, and `value_gradient_and_hessian!!` now accept a `context=` keyword that overrides, for a single call, the `context` frozen at `prepare` — without re-preparing (#167). The default (`context=nothing`) leaves the hot path unchanged; a `Tuple` override must match the frozen context's element types and shapes — a type mismatch (or a non-`Tuple` override) throws an `ArgumentError` — and reuses the type-keyed cache. Backends that bake context into their prepared state throw an `ArgumentError` for a non-empty-input override (re-`prepare` instead): compiled-tape ReverseDiff (`AutoReverseDiff(; compile=true)`) on all three entry points, and Mooncake's Hessian; Mooncake Jacobian preps are context-free by construction, so only an empty override validates there. Empty input runs no derivative machinery, so no backend rejects an override there (it is still validated). - -Fixed a Mooncake reverse-mode correctness bug ([Mooncake issue #1238](https://github.com/chalk-lab/Mooncake.jl/issues/1238)): when a problem function closed over differentiable data (e.g. a model capturing observed data reached through an in-place solve), a reused prepared cache accumulated that captured data's cotangent and returned an incorrect gradient after the first evaluation. The Mooncake extension now closes the function and its context into a `NoTangent` target, so the captured data carries no cotangent and reuse is correct. +## 0.16.0 Moved `VarNamedTuple` and its structured-array support from DynamicPPL into AbstractPPL so probabilistic programming frontends can share a lossless, `VarName`-keyed parameter @@ -10,6 +6,12 @@ container. The exported API includes `VarNamedTuple`, `@vnt`, `subset`, and the mapping, densification, and skeleton helpers. Template-based insertion remains public as `templated_setindex!!`. ComponentArrays integration is provided by a package extension. +## 0.15.4 + +`value_and_gradient!!`, `value_and_jacobian!!`, and `value_gradient_and_hessian!!` now accept a `context=` keyword that overrides, for a single call, the `context` frozen at `prepare` — without re-preparing (#167). The default (`context=nothing`) leaves the hot path unchanged; a `Tuple` override must match the frozen context's element types and shapes — a type mismatch (or a non-`Tuple` override) throws an `ArgumentError` — and reuses the type-keyed cache. Backends that bake context into their prepared state throw an `ArgumentError` for a non-empty-input override (re-`prepare` instead): compiled-tape ReverseDiff (`AutoReverseDiff(; compile=true)`) on all three entry points, and Mooncake's Hessian; Mooncake Jacobian preps are context-free by construction, so only an empty override validates there. Empty input runs no derivative machinery, so no backend rejects an override there (it is still validated). + +Fixed a Mooncake reverse-mode correctness bug ([Mooncake issue #1238](https://github.com/chalk-lab/Mooncake.jl/issues/1238)): when a problem function closed over differentiable data (e.g. a model capturing observed data reached through an in-place solve), a reused prepared cache accumulated that captured data's cotangent and returned an incorrect gradient after the first evaluation. The Mooncake extension now closes the function and its context into a `NoTangent` target, so the captured data carries no cotangent and reuse is correct. + ## 0.15.3 Added the `of` type system: a self-contained, declarative way to specify the shape, element type, and support of model variables. Construct specifications with the exported `of` function or the exported `@of` macro: diff --git a/Project.toml b/Project.toml index f99d7d62..ca07c675 100644 --- a/Project.toml +++ b/Project.toml @@ -3,7 +3,7 @@ uuid = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" keywords = ["probabilistic programming"] license = "MIT" desc = "Common interfaces for probabilistic programming" -version = "0.15.4" +version = "0.16.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" diff --git a/src/varnamedtuple/getset.jl b/src/varnamedtuple/getset.jl index 356039a9..43fd24ac 100644 --- a/src/varnamedtuple/getset.jl +++ b/src/varnamedtuple/getset.jl @@ -436,7 +436,7 @@ function make_leaf_multiindex(value, coptic::AbstractPPL.Index, template) pa_eltype = if sub_value isa AbstractArray || sub_value isa PartialArray eltype(sub_value) else - idx_size_type = Dims{_ndims(template, coptic.ix...;coptic.kw...)} + idx_size_type = Dims{_ndims(template, coptic.ix...; coptic.kw...)} ArrayLikeBlock{typeof(sub_value),typeof(coptic.ix),typeof(coptic.kw),idx_size_type} end diff --git a/src/varnamedtuple/vnt.jl b/src/varnamedtuple/vnt.jl index 22fae669..18d1aac7 100644 --- a/src/varnamedtuple/vnt.jl +++ b/src/varnamedtuple/vnt.jl @@ -234,16 +234,19 @@ end if !_has_partial_array(ValType) return :(return false) end - push!(exs, quote - val = vnt.data.$name - if val isa VarNamedTuple || val isa PartialArray - if !Base.isempty(val) + push!( + exs, + quote + val = vnt.data.$name + if val isa VarNamedTuple || val isa PartialArray + if !Base.isempty(val) + return false + end + else return false end - else - return false - end - end) + end, + ) end push!(exs, :(return true)) return Expr(:block, exs...) diff --git a/test/varname/hasvalue.jl b/test/varname/hasvalue.jl index 085cfc90..c6cc7d8f 100644 --- a/test/varname/hasvalue.jl +++ b/test/varname/hasvalue.jl @@ -34,11 +34,11 @@ using Test @test canview(@opticof(_[1:2]), x) @test canview(@opticof(_[:]), x) @test !canview(@opticof(_[4]), x) - @test canview(@opticof(_[i = 1]), x) + @test canview(@opticof(_[i=1]), x) # For some weird reason DimData does not error on these two but just warns that # there's no index j! - @test canview(@opticof(_[j = 2]), x) - @test canview(@opticof(_[i = 1, j = 2]), x) + @test canview(@opticof(_[j=2]), x) + @test canview(@opticof(_[i=1, j=2]), x) end @testset "Dict" begin @@ -246,14 +246,14 @@ end @test getvalue(x, @varname(a[1, 2])) == x.a[1, 2] @test hasvalue(x, @varname(a[:])) @test getvalue(x, @varname(a[:])) == x.a[:] - @test canview(@opticof(_[i = 1]), x.a) - @test hasvalue(x, @varname(a[i = 1])) - @test getvalue(x, @varname(a[i = 1])) == x.a[i = 1] - @test canview(@opticof(_[i = 1, j = 2]), x.a) - @test hasvalue(x, @varname(a[i = 1, j = 2])) - @test getvalue(x, @varname(a[i = 1, j = 2])) == x.a[i = 1, j = 2] - @test hasvalue(x, @varname(a[i = DD.Not(1)])) - @test getvalue(x, @varname(a[i = DD.Not(1)])) == x.a[i = DD.Not(1)] + @test canview(@opticof(_[i=1]), x.a) + @test hasvalue(x, @varname(a[i=1])) + @test getvalue(x, @varname(a[i=1])) == x.a[i=1] + @test canview(@opticof(_[i=1, j=2]), x.a) + @test hasvalue(x, @varname(a[i=1, j=2])) + @test getvalue(x, @varname(a[i=1, j=2])) == x.a[i=1, j=2] + @test hasvalue(x, @varname(a[i=DD.Not(1)])) + @test getvalue(x, @varname(a[i=DD.Not(1)])) == x.a[i=DD.Not(1)] y = (; b=DD.DimArray(randn(2, 3), (DD.X, DD.Y))) @test hasvalue(y, @varname(b)) From 68b04e8a162690db09cbc459d7b9719913157485 Mon Sep 17 00:00:00 2001 From: Xianda Sun Date: Mon, 13 Jul 2026 09:09:52 +0100 Subject: [PATCH 3/3] Note additivity and the DynamicPPL export overlap in the 0.16.0 entry --- HISTORY.md | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/HISTORY.md b/HISTORY.md index 6933204d..9f54e85c 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -4,7 +4,18 @@ Moved `VarNamedTuple` and its structured-array support from DynamicPPL into Abst probabilistic programming frontends can share a lossless, `VarName`-keyed parameter container. The exported API includes `VarNamedTuple`, `@vnt`, `subset`, and the existing mapping, densification, and skeleton helpers. Template-based insertion remains public as -`templated_setindex!!`. ComponentArrays integration is provided by a package extension. +`templated_setindex!!`. ComponentArrays integration is provided by a package extension, and +the `Distribution`-aware `hasvalue`/`getvalue` methods live in the existing Distributions +package extension. + +This release is purely additive over 0.15 — nothing was removed or changed in the existing +public API, so downstream packages can widen their compat bound without code changes. +However, the ten newly exported names (`VarNamedTuple`, `@vnt`, `map_pairs!!`, +`map_values!!`, `apply!!`, `densify!!`, `skeleton`, `subset`, `NoTemplate`, `SkipTemplate`) +are currently also exported by DynamicPPL for its own copy of the same functionality. Code +that does `using DynamicPPL, AbstractPPL` unqualified will therefore see ambiguous bindings +for these names until DynamicPPL migrates to re-export AbstractPPL's implementation; +qualify uses in the meantime. ## 0.15.4