Implement prefix / unprefix

Author: penelopeysm

src/AbstractPPL.jl

@@ -14,7 +14,9 @@ export VarName,
     index_to_dict,
     dict_to_index,
     varname_to_string,
-    string_to_varname
+    string_to_varname,
+    prefix,
+    unprefix
 
 # Abstract model functions
 export AbstractProbabilisticProgram,

src/varname.jl

@@ -766,6 +766,8 @@ function vsym(expr::Expr)
     end
 end
 
+### Serialisation to JSON / string
+
 # String constants for each index type that we support serialisation /
 # deserialisation of
 const _BASE_INTEGER_TYPE = "Base.Integer"
@@ -936,3 +938,152 @@ 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))
+
+### Prefixing and unprefixing
+
+"""
+    _strip_identity(optic)
+
+Remove an inner layer of the identity lens from a composed optic.
+"""
+_strip_identity(o::Base.ComposedFunction{Outer,typeof(identity)}) where {Outer} = o.outer
+_strip_identity(o::Base.ComposedFunction) = o
+_strip_identity(o::Accessors.PropertyLens) = o
+_strip_identity(o::Accessors.IndexLens) = o
+_strip_identity(o::typeof(identity)) = o
+
+"""
+    _inner(optic)
+
+Get the innermost (non-identity) layer of an optic.
+
+```jldoctest; setup=:(using Accessors)
+julia> AbstractPPL._inner(Accessors.@o _.a.b.c)
+(@o _.a)
+
+julia> AbstractPPL._inner(Accessors.@o _[1][2][3])
+(@o _[1])
+
+julia> AbstractPPL._inner(Accessors.@o _)
+identity (generic function with 1 method)
+```
+"""
+_inner(o::Base.ComposedFunction{Outer,Inner}) where {Outer,Inner} = o.inner
+function _inner(o::Base.ComposedFunction{Outer,typeof(identity)}) where {Outer}
+    return _strip_identity(o.outer)
+end
+_inner(o::Accessors.PropertyLens) = o
+_inner(o::Accessors.IndexLens) = o
+_inner(o::typeof(identity)) = o
+
+"""
+    _outer(optic)
+
+Get the outer layer of an optic.
+
+```jldoctest; setup=:(using Accessors)
+julia> AbstractPPL._outer(Accessors.@o _.a.b.c)
+(@o _.b.c)
+
+julia> AbstractPPL._outer(Accessors.@o _[1][2][3])
+(@o _[2][3])
+
+julia> AbstractPPL._outer(Accessors.@o _.a)
+identity (generic function with 1 method)
+
+julia> AbstractPPL._outer(Accessors.@o _[1])
+identity (generic function with 1 method)
+
+julia> AbstractPPL._outer(Accessors.@o _)
+identity (generic function with 1 method)
+```
+"""
+_outer(o::Base.ComposedFunction{Outer,Inner}) where {Outer,Inner} = _strip_identity(o.outer)
+_outer(::Accessors.PropertyLens) = identity
+_outer(::Accessors.IndexLens) = identity
+_outer(::typeof(identity)) = identity
+
+"""
+    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::Base.ComposedFunction{Outer,typeof(identity)}) where {Outer}
+    return optic_to_vn(o.outer)
+end
+function optic_to_vn(
+    o::Base.ComposedFunction{Outer,Accessors.PropertyLens{sym}}
+) where {Outer,sym}
+    return VarName{sym}(o.outer)
+end
+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)
+    # strip one layer of the optic and check for equality
+    inner = _inner(optic)
+    inner_prefix = _inner(optic_prefix)
+    if inner != inner_prefix
+        msg = "could not remove prefix $(optic_prefix) from optic $(optic)"
+        throw(ArgumentError(msg))
+    end
+    # recurse
+    return unprefix_optic(_outer(optic), _outer(optic_prefix))
+end
+
+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
+
+function prefix(vn::VarName{sym_vn}, prefix::VarName{sym_prefix}) where {sym_vn,sym_prefix}
+    optic_vn = getoptic(vn)
+    optic_prefix = getoptic(prefix)
+    # Special case `identity` to avoid having ComposedFunctions with identity
+    if optic_vn == identity
+        new_inner_optic_vn = PropertyLens{sym_vn}()
+    else
+        new_inner_optic_vn = optic_vn ∘ PropertyLens{sym_vn}()
+    end
+    if optic_prefix == identity
+        new_optic_vn = new_inner_optic_vn
+    else
+        new_optic_vn = new_inner_optic_vn ∘ optic_prefix
+    end
+    return VarName{sym_prefix}(new_optic_vn)
+end

test/varname.jl

@@ -233,4 +233,18 @@ end
         # Serialisation should now work
         @test string_to_varname(varname_to_string(vn)) == vn
     end
+
+    @testset "prefix and unprefix" begin
+        @test prefix(@varname(y), @varname(x)) == @varname(x.y)
+        @test prefix(@varname(y), @varname(x[1])) == @varname(x[1].y)
+        @test prefix(@varname(y), @varname(x.a)) == @varname(x.a.y)
+        @test prefix(@varname(y[1]), @varname(x)) == @varname(x.y[1])
+        @test prefix(@varname(y.a), @varname(x)) == @varname(x.y.a)
+
+        @test unprefix(@varname(x.y[1]), @varname(x)) == @varname(y[1])
+        @test unprefix(@varname(x[1].y), @varname(x[1])) == @varname(y)
+        @test unprefix(@varname(x.a.y), @varname(x.a)) == @varname(y)
+        @test unprefix(@varname(x.y.a), @varname(x)) == @varname(y.a)
+        @test_throws ArgumentError unprefix(@varname(x.y.a), @varname(n))
+    end
 end