Split up varname.jl into more manageable bits

.github/workflows/CI.yml

@@ -28,7 +28,7 @@ jobs:
             os: macos-latest
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
 
       - uses: julia-actions/setup-julia@v2
         with:
@@ -49,8 +49,3 @@ jobs:
           files: lcov.info
           token: ${{ secrets.CODECOV_TOKEN }}
           fail_ci_if_error: true
-
-      - uses: coverallsapp/github-action@v2
-        with:
-          github-token: ${{ secrets.GITHUB_TOKEN }}
-          path-to-lcov: lcov.info

README.md

@@ -2,7 +2,6 @@
 
 [![CI](https://github.com/TuringLang/AbstractPPL.jl/actions/workflows/CI.yml/badge.svg?branch=main)](https://github.com/TuringLang/AbstractPPL.jl/actions/workflows/CI.yml?query=branch%3Amain)
 [![IntegrationTest](https://github.com/TuringLang/AbstractPPL.jl/actions/workflows/IntegrationTest.yml/badge.svg?branch=main)](https://github.com/TuringLang/AbstractPPL.jl/actions/workflows/IntegrationTest.yml?query=branch%3Amain)
-[![Coverage Status](https://coveralls.io/repos/github/TuringLang/AbstractPPL.jl/badge.svg?branch=main)](https://coveralls.io/github/TuringLang/AbstractPPL.jl?branch=main)
 [![Codecov](https://codecov.io/gh/TuringLang/AbstractPPL.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/TuringLang/AbstractPPL.jl)
 
 A light-weight package to factor out interfaces and associated APIs for modelling languages for

src/AbstractPPL.jl

@@ -29,11 +29,14 @@ export AbstractProbabilisticProgram,
 # Abstract traces
 export AbstractModelTrace
 
-include("varname.jl")
 include("abstractmodeltrace.jl")
 include("abstractprobprog.jl")
 include("evaluate.jl")
-include("varname_leaves.jl")
-include("hasvalue.jl")
+include("varname/varname.jl")
+include("varname/subsumes.jl")
+include("varname/hasvalue.jl")
+include("varname/leaves.jl")
+include("varname/prefix.jl")
+include("varname/serialize.jl")
 
 end # module

src/varname/prefix.jl

@@ -0,0 +1,116 @@
+### Functionality for prefixing and unprefixing VarNames.
+
+"""
+    optic_to_vn(optic)
+
+Convert an Accessors optic to a VarName. This is best explained through
+examples.
+
+```jldoctest; setup=:(using Accessors)
+julia> AbstractPPL.optic_to_vn(Accessors.@o _.a)
+a
+
+julia> AbstractPPL.optic_to_vn(Accessors.@o _.a.b)
+a.b
+
+julia> AbstractPPL.optic_to_vn(Accessors.@o _.a[1])
+a[1]
+```
+
+The outermost layer of the optic (technically, what Accessors.jl calls the
+'innermost') must be a `PropertyLens`, or else it will fail. This is because a
+VarName needs to have a symbol.
+
+```jldoctest; setup=:(using Accessors)
+julia> AbstractPPL.optic_to_vn(Accessors.@o _[1])
+ERROR: ArgumentError: optic_to_vn: could not convert optic `(@o _[1])` to a VarName
+[...]
+```
+"""
+function optic_to_vn(::Accessors.PropertyLens{sym}) where {sym}
+    return VarName{sym}()
+end
+function optic_to_vn(
+    o::ComposedFunction{Outer,Accessors.PropertyLens{sym}}
+) where {Outer,sym}
+    return VarName{sym}(o.outer)
+end
+optic_to_vn(o::ComposedFunction) = optic_to_vn(normalise(o))
+function optic_to_vn(@nospecialize(o))
+    msg = "optic_to_vn: could not convert optic `$o` to a VarName"
+    throw(ArgumentError(msg))
+end
+
+unprefix_optic(o, ::typeof(identity)) = o  # Base case
+function unprefix_optic(optic, optic_prefix)
+    # Technically `unprefix_optic` only receives optics that were part of
+    # VarNames, so the optics should already be normalised (in the inner
+    # constructor of the VarName). However I guess it doesn't hurt to do it
+    # again to be safe.
+    optic = normalise(optic)
+    optic_prefix = normalise(optic_prefix)
+    # strip one layer of the optic and check for equality
+    head = _head(optic)
+    head_prefix = _head(optic_prefix)
+    if head != head_prefix
+        msg = "could not remove prefix $(optic_prefix) from optic $(optic)"
+        throw(ArgumentError(msg))
+    end
+    # recurse
+    return unprefix_optic(_tail(optic), _tail(optic_prefix))
+end
+
+"""
+    unprefix(vn::VarName, prefix::VarName)
+
+Remove a prefix from a VarName.
+
+```jldoctest
+julia> AbstractPPL.unprefix(@varname(y.x), @varname(y))
+x
+
+julia> AbstractPPL.unprefix(@varname(y.x.a), @varname(y))
+x.a
+
+julia> AbstractPPL.unprefix(@varname(y[1].x), @varname(y[1]))
+x
+
+julia> AbstractPPL.unprefix(@varname(y), @varname(n))
+ERROR: ArgumentError: could not remove prefix n from VarName y
+[...]
+```
+"""
+function unprefix(
+    vn::VarName{sym_vn}, prefix::VarName{sym_prefix}
+) where {sym_vn,sym_prefix}
+    if sym_vn != sym_prefix
+        msg = "could not remove prefix $(prefix) from VarName $(vn)"
+        throw(ArgumentError(msg))
+    end
+    optic_vn = getoptic(vn)
+    optic_prefix = getoptic(prefix)
+    return optic_to_vn(unprefix_optic(optic_vn, optic_prefix))
+end
+
+"""
+    prefix(vn::VarName, prefix::VarName)
+
+Add a prefix to a VarName.
+
+```jldoctest
+julia> AbstractPPL.prefix(@varname(x), @varname(y))
+y.x
+
+julia> AbstractPPL.prefix(@varname(x.a), @varname(y))
+y.x.a
+
+julia> AbstractPPL.prefix(@varname(x.a), @varname(y[1]))
+y[1].x.a
+```
+"""
+function prefix(vn::VarName{sym_vn}, prefix::VarName{sym_prefix}) where {sym_vn,sym_prefix}
+    optic_vn = getoptic(vn)
+    optic_prefix = getoptic(prefix)
+    new_optic_vn = optic_vn ∘ PropertyLens{sym_vn}() ∘ optic_prefix
+    return VarName{sym_prefix}(new_optic_vn)
+end

src/varname/serialize.jl

@@ -0,0 +1,174 @@
+### Serialisation to JSON / string
+
+using JSON: JSON
+
+# String constants for each index type that we support serialisation /
+# deserialisation of
+const _BASE_INTEGER_TYPE = "Base.Integer"
+const _BASE_VECTOR_TYPE = "Base.Vector"
+const _BASE_UNITRANGE_TYPE = "Base.UnitRange"
+const _BASE_STEPRANGE_TYPE = "Base.StepRange"
+const _BASE_ONETO_TYPE = "Base.OneTo"
+const _BASE_COLON_TYPE = "Base.Colon"
+const _CONCRETIZED_SLICE_TYPE = "AbstractPPL.ConcretizedSlice"
+const _BASE_TUPLE_TYPE = "Base.Tuple"
+
+"""
+    index_to_dict(::Integer)
+    index_to_dict(::AbstractVector{Int})
+    index_to_dict(::UnitRange)
+    index_to_dict(::StepRange)
+    index_to_dict(::Colon)
+    index_to_dict(::ConcretizedSlice{T, Base.OneTo{I}}) where {T, I}
+    index_to_dict(::Tuple)
+
+Convert an index `i` to a dictionary representation.
+"""
+index_to_dict(i::Integer) = Dict("type" => _BASE_INTEGER_TYPE, "value" => i)
+index_to_dict(v::Vector{Int}) = Dict("type" => _BASE_VECTOR_TYPE, "values" => v)
+function index_to_dict(r::UnitRange)
+    return Dict("type" => _BASE_UNITRANGE_TYPE, "start" => r.start, "stop" => r.stop)
+end
+function index_to_dict(r::StepRange)
+    return Dict(
+        "type" => _BASE_STEPRANGE_TYPE,
+        "start" => r.start,
+        "stop" => r.stop,
+        "step" => r.step,
+    )
+end
+function index_to_dict(r::Base.OneTo{I}) where {I}
+    return Dict("type" => _BASE_ONETO_TYPE, "stop" => r.stop)
+end
+index_to_dict(::Colon) = Dict("type" => _BASE_COLON_TYPE)
+function index_to_dict(s::ConcretizedSlice{T,R}) where {T,R}
+    return Dict("type" => _CONCRETIZED_SLICE_TYPE, "range" => index_to_dict(s.range))
+end
+function index_to_dict(t::Tuple)
+    return Dict("type" => _BASE_TUPLE_TYPE, "values" => map(index_to_dict, t))
+end
+
+"""
+    dict_to_index(dict)
+    dict_to_index(symbol_val, dict)
+
+Convert a dictionary representation of an index `dict` to an index.
+
+Users can extend the functionality of `dict_to_index` (and hence `VarName`
+de/serialisation) by extending this method along with [`index_to_dict`](@ref).
+Specifically, suppose you have a custom index type `MyIndexType` and you want
+to be able to de/serialise a `VarName` containing this index type. You should
+then implement the following two methods:
+
+1. `AbstractPPL.index_to_dict(i::MyModule.MyIndexType)` should return a
+   dictionary representation of the index `i`. This dictionary must contain the
+   key `"type"`, and the corresponding value must be a string that uniquely
+   identifies the index type. Generally, it makes sense to use the name of the
+   type (perhaps prefixed with module qualifiers) as this value to avoid
+   clashes. The remainder of the dictionary can have any structure you like.
+
+2. Suppose the value of `index_to_dict(i)["type"]` is `"MyModule.MyIndexType"`.
+   You should then implement the corresponding method
+   `AbstractPPL.dict_to_index(::Val{Symbol("MyModule.MyIndexType")}, dict)`,
+   which should take the dictionary representation as the second argument and
+   return the original `MyIndexType` object.
+
+To see an example of this in action, you can look in the the AbstractPPL test
+suite, which contains a test for serialising OffsetArrays.
+"""
+function dict_to_index(dict)
+    t = dict["type"]
+    if t == _BASE_INTEGER_TYPE
+        return dict["value"]
+    elseif t == _BASE_VECTOR_TYPE
+        return collect(Int, dict["values"])
+    elseif t == _BASE_UNITRANGE_TYPE
+        return dict["start"]:dict["stop"]
+    elseif t == _BASE_STEPRANGE_TYPE
+        return dict["start"]:dict["step"]:dict["stop"]
+    elseif t == _BASE_ONETO_TYPE
+        return Base.OneTo(dict["stop"])
+    elseif t == _BASE_COLON_TYPE
+        return Colon()
+    elseif t == _CONCRETIZED_SLICE_TYPE
+        return ConcretizedSlice(Base.Slice(dict_to_index(dict["range"])))
+    elseif t == _BASE_TUPLE_TYPE
+        return tuple(map(dict_to_index, dict["values"])...)
+    else
+        # Will error if the method is not defined, but this hook allows users
+        # to extend this function
+        return dict_to_index(Val(Symbol(t)), dict)
+    end
+end
+
+optic_to_dict(::typeof(identity)) = Dict("type" => "identity")
+function optic_to_dict(::PropertyLens{sym}) where {sym}
+    return Dict("type" => "property", "field" => String(sym))
+end
+optic_to_dict(i::IndexLens) = Dict("type" => "index", "indices" => index_to_dict(i.indices))
+function optic_to_dict(c::ComposedFunction)
+    return Dict(
+        "type" => "composed",
+        "outer" => optic_to_dict(c.outer),
+        "inner" => optic_to_dict(c.inner),
+    )
+end
+
+function dict_to_optic(dict)
+    if dict["type"] == "identity"
+        return identity
+    elseif dict["type"] == "index"
+        return IndexLens(dict_to_index(dict["indices"]))
+    elseif dict["type"] == "property"
+        return PropertyLens{Symbol(dict["field"])}()
+    elseif dict["type"] == "composed"
+        return dict_to_optic(dict["outer"]) ∘ dict_to_optic(dict["inner"])
+    else
+        error("Unknown optic type: $(dict["type"])")
+    end
+end
+
+function varname_to_dict(vn::VarName)
+    return Dict("sym" => getsym(vn), "optic" => optic_to_dict(getoptic(vn)))
+end
+
+function dict_to_varname(dict::Dict{<:AbstractString,Any})
+    return VarName{Symbol(dict["sym"])}(dict_to_optic(dict["optic"]))
+end
+
+"""
+    varname_to_string(vn::VarName)
+
+Convert a `VarName` as a string, via an intermediate dictionary. This differs
+from `string(vn)` in that concretised slices are faithfully represented (rather
+than being pretty-printed as colons).
+
+For `VarName`s which index into an array, this function will only work if the
+indices can be serialised. This is true for all standard Julia index types, but
+if you are using custom index types, you will need to implement the
+`index_to_dict` and `dict_to_index` methods for those types. See the
+documentation of [`dict_to_index`](@ref) for instructions on how to do this.
+
+```jldoctest
+julia> varname_to_string(@varname(x))
+"{\\"optic\\":{\\"type\\":\\"identity\\"},\\"sym\\":\\"x\\"}"
+
+julia> varname_to_string(@varname(x.a))
+"{\\"optic\\":{\\"field\\":\\"a\\",\\"type\\":\\"property\\"},\\"sym\\":\\"x\\"}"
+
+julia> y = ones(2); varname_to_string(@varname(y[:]))
+"{\\"optic\\":{\\"indices\\":{\\"values\\":[{\\"type\\":\\"Base.Colon\\"}],\\"type\\":\\"Base.Tuple\\"},\\"type\\":\\"index\\"},\\"sym\\":\\"y\\"}"
+
+julia> y = ones(2); varname_to_string(@varname(y[:], true))
+"{\\"optic\\":{\\"indices\\":{\\"values\\":[{\\"range\\":{\\"stop\\":2,\\"type\\":\\"Base.OneTo\\"},\\"type\\":\\"AbstractPPL.ConcretizedSlice\\"}],\\"type\\":\\"Base.Tuple\\"},\\"type\\":\\"index\\"},\\"sym\\":\\"y\\"}"
+```
+"""
+varname_to_string(vn::VarName) = JSON.json(varname_to_dict(vn))
+
+"""
+    string_to_varname(str::AbstractString)
+
+Convert a string representation of a `VarName` back to a `VarName`. The string
+should have been generated by `varname_to_string`.
+"""
+string_to_varname(str::AbstractString) = dict_to_varname(JSON.parse(str))