Using Setfield.jl's lenses to handle indexing

Project.toml

@@ -7,6 +7,8 @@ version = "0.2.0"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
+Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [compat]
 AbstractMCMC = "2, 3"

src/varname.jl

@@ -1,3 +1,6 @@
+using Setfield
+import Setfield: PropertyLens, ComposedLens, IdentityLens, IndexLens, DynamicIndexLens
+
 """
     VarName{sym}(indexing::Tuple=())
 
@@ -26,10 +29,10 @@ julia> @varname x[:, 1][1+1]
 x[:,1][2]
 ```
 """
-struct VarName{sym, T<:Tuple}
+struct VarName{sym, T}
     indexing::T
 
-    VarName{sym}(indexing::Tuple=()) where {sym} = new{sym,typeof(indexing)}(indexing)
+    VarName{sym}(indexing=()) where {sym} = new{sym,typeof(indexing)}(indexing)
 end
 
 """
@@ -45,7 +48,7 @@ julia> VarName(@varname(x[1][2:3]))
 x
 ```
 """
-function VarName(vn::VarName, indexing::Tuple = ())
+function VarName(vn::VarName, indexing = ())
     return VarName{getsym(vn)}(indexing)
 end
 
@@ -89,7 +92,7 @@ getindexing(vn::VarName) = vn.indexing
 Base.hash(vn::VarName, h::UInt) = hash((getsym(vn), getindexing(vn)), h)
 Base.:(==)(x::VarName, y::VarName) = getsym(x) == getsym(y) && getindexing(x) == getindexing(y)
 
-function Base.show(io::IO, vn::VarName)
+function Base.show(io::IO, vn::VarName{<:Any, <:Tuple})
     print(io, getsym(vn))
     for indices in getindexing(vn)
         print(io, "[")
@@ -98,6 +101,15 @@ function Base.show(io::IO, vn::VarName)
     end
 end
 
+function Base.show(io::IO, vn::VarName{<:Any, <:Lens})
+    print(io, getsym(vn))
+    Setfield.print_application(io, vn.indexing)
+end
+
+# TODO: Should this really go here?
+Setfield.print_application(io::IO, l::IndexLens) = print(io, "[", join(map(prettify_index, l.indices), ", "), "]")
+Setfield.print_application(io::IO, l::DynamicIndexLens) = print(io, l, "(_)")
+
 prettify_index(x) = string(x)
 prettify_index(::Colon) = ":"
 
@@ -214,7 +226,65 @@ _issubrange(i::ConcreteIndex, j::ConcreteIndex) = issubset(i, j)
 _issubrange(i::Union{ConcreteIndex, Colon}, j::Colon) = true
 _issubrange(i::Colon, j::ConcreteIndex) = true
 
+# Idea behind `subsumes` for `Lens` is that we traverse the two lenses in parallel,
+# checking `subsumes` for every level. This for example means that if we are comparing
+# `PropertyLens{:a}` and `PropertyLens{:b}` we immediately know that they do not subsume
+# each other since at the same level/depth they access different properties.
+subsumes(t::ComposedLens, u::ComposedLens) = subsumes(t.outer, u.outer) && subsumes(t.inner, u.inner)
+
+# If `t` is still a composed lens, then there is no way it can subsume `u` since `u` is a
+# leaf of the "lens-tree".
+subsumes(t::ComposedLens, u::PropertyLens) = false
+# Here we need to check if `u.outer` (i.e. the next lens to be applied from `u`) is
+# subsumed by `t`, since this would mean that the rest of the composition is also subsumed
+# by `t`.
+subsumes(t::PropertyLens, u::ComposedLens) = subsumes(t, u.outer)
+
+# For `PropertyLens` either they have the same `name` and thus they are indeed the same.
+subsumes(t::PropertyLens{name}, u::PropertyLens{name}) where {name} = true
+# Otherwise they represent different properties, and thus are not the same.
+subsumes(t::PropertyLens, u::PropertyLens) = false
+
+# Indices subsumes if they are subindices, i.e. we just call `_issubindex`.
+# FIXME: Does not support `DynamicIndexLens`.
+# FIXME: Does not correctly handle cases such as `subsumes(x, x[:])`
+#        (but neither did old implementation).
+subsumes(t::IndexLens, u::IndexLens) = _issubindex(t.indices, u.indices)
+
+"""
+    concretize(l::Lens, x)
 
+Return `l` instantiated on `x`, i.e. any runtime information evaluated using `x`.
+"""
+
+concretize(I::Lens, x) = I
+concretize(I::DynamicIndexLens, x) = IndexLens(I.f(x))
+function concretize(I::ComposedLens, x)
+    x_inner = get(x, I.outer)
+    return ComposedLens(concretize(I.outer, x), concretize(I.inner, x_inner))
+end
+"""
+    concretize(vn::VarName, x)
+
+Return `vn` instantiated on `x`, i.e. any runtime information evaluated using `x`.
+
+# Examples
+```jldoctest
+julia> x = (a = [1.0 2.0;], );
+
+julia> vn = @varname(x.a[1, :])
+x.a[1, :]
+
+julia> AbstractPPL.concretize(vn, x)
+x.a[1, :]
+
+julia> vn = @varname(x.a[1, end][:]);
+
+julia> AbstractPPL.concretize(vn, x)
+x.a[1, 2][:]
+```
+"""
+concretize(vn::VarName, x) = VarName(vn, concretize(vn.indexing, x))
 
 """
     @varname(expr)
@@ -231,30 +301,43 @@ julia> @varname(x).indexing
 ()
 
 julia> @varname(x[1]).indexing
-((1,),)
+(@lens _[1])
 
 julia> @varname(x[:, 1]).indexing
-((Colon(), 1),)
+(@lens _[:, 1])
 
 julia> @varname(x[:, 1][2]).indexing
-((Colon(), 1), (2,))
+(@lens _[:, 1][2])
 
 julia> @varname(x[1,2][1+5][45][3]).indexing
-((1, 2), (6,), (45,), (3,))
+(@lens _[1, 2][6][45][3])
 ```
 
 !!! compat "Julia 1.5"
     Using `begin` in an indexing expression to refer to the first index requires at least
     Julia 1.5.
 """
 macro varname(expr::Union{Expr, Symbol})
-    return esc(varname(expr))
+    return varname(expr)
 end
 
 varname(sym::Symbol) = :($(AbstractPPL.VarName){$(QuoteNode(sym))}())
 function varname(expr::Expr)
-    if Meta.isexpr(expr, :ref)
-        sym, inds = vsym(expr), vinds(expr)
+    if Meta.isexpr(expr, :ref) || Meta.isexpr(expr, :.)
+        sym = vsym(expr)
+
+        # Need to recursively unwrap until we reach the outer-most variable.
+        # TODO: implement as recursion?
+        curexpr = expr
+        while !(curexpr.args[1] isa Symbol)
+            curexpr = curexpr.args[1]
+        end
+
+        # Then we replace the variable with `_`, to get an expression we can
+        # use `lensmacro` on.
+        curexpr.args[1] = :_
+        inds = Setfield.lensmacro(identity, expr)
+
         return :($(AbstractPPL.VarName){$(QuoteNode(sym))}($inds))
     else
         error("Malformed variable name $(expr)!")
@@ -294,7 +377,7 @@ function vsym end
 
 vsym(expr::Symbol) = expr
 function vsym(expr::Expr)
-    if Meta.isexpr(expr, :ref)
+    if Meta.isexpr(expr, :ref) || Meta.isexpr(expr, :.)
         return vsym(expr.args[1])
     else
         error("Malformed variable name $(expr)!")