From 69e39248953db63c420817790e5fffda8f2d5703 Mon Sep 17 00:00:00 2001
From: Philipp Gabler <phipsgabler@users.noreply.github.com>
Date: Tue, 7 Feb 2023 20:47:08 +0000
Subject: [PATCH] Allow name and property interpolation in varnames (#54)

Implements #46.

Most of this is dependent on my [PR in Setfield.jl](https://github.com/jw3126/Setfield.jl/pull/168) to allow interpolation of properties (so it works only if you `dev` that branch).

It's a bit of a hack though, since `parse_object_lens` behaves a bit weirdly when the innermost value of a lens chain is an interpolation:

```julia
julia> Setfield.parse_obj_lens(@q $name.a)
(:($(Expr(:escape, :_))), :((Setfield.compose)($(Expr(:escape, :name)), (Setfield.PropertyLens){:a}())))

julia> Setfield.parse_obj_lens(@q x.a)
(:($(Expr(:escape, :x))), :((Setfield.compose)((Setfield.PropertyLens){:a}())))
```

I'd be really nice to always get the latter form.  (`@q` is just an ad-hoc macro to construct expressions preserving the `:$`.)

@jw3126, could you maybe look at this and suggest any better alternatives? Either how to implement this nicely in AbstractPPL.jl, or improvements I can add to my PR.

Co-authored-by: Xianda <sunxdt@gmail.com>
---
 Project.toml   |  2 +-
 src/varname.jl | 41 +++++++++++++++++++++++++++++++++++++++--
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/Project.toml b/Project.toml
index 8da7c9f..ba18628 100644
--- a/Project.toml
+++ b/Project.toml
@@ -15,5 +15,5 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 [compat]
 AbstractMCMC = "2, 3, 4"
 DensityInterface = "0.4"
-Setfield = "0.8.1, 1"
+Setfield = "0.8.2, 1"
 julia = "~1.6.6, 1.7.3"
diff --git a/src/varname.jl b/src/varname.jl
index c6fde32..a0267e3 100644
--- a/src/varname.jl
+++ b/src/varname.jl
@@ -515,6 +515,7 @@ concretized as `VarName` only supports non-dynamic indexing as determined by
 [`is_static_index`](@ref). See examples below.
 
 ## Examples
+
 ### Dynamic indexing
 ```jldoctest
 julia> x = (a = [1.0 2.0; 3.0 4.0; 5.0 6.0], );
@@ -573,6 +574,28 @@ julia> x = (a = [(b = rand(2), )], ); getlens(@varname(x.a[1].b[end], true))
 (@lens _.a[1].b[2])
 ```
 
+Interpolation can be used for names (the base name as well as property names).  Variables within
+indices are always evaluated in the calling scope, in the same manner as `Setfield` does:
+
+```jldoctest
+julia> name, i = :a, 10;
+
+julia> @varname(x.\$name[i, i+1])
+x.a[10,11]
+
+julia> @varname(\$name)
+a
+
+julia> @varname(\$name[1])
+a[1]
+
+julia> @varname(\$name.x[1])
+a.x[1]
+
+julia> @varname(b.\$name.x[1])
+b.a.x[1]
+```
+
 !!! compat "Julia 1.5"
     Using `begin` in an indexing expression to refer to the first index requires at least
     Julia 1.5.
@@ -591,17 +614,31 @@ function varname(expr::Expr, concretize=Setfield.need_dynamic_lens(expr))
         # to call `QuoteNode` on it.
         sym = drop_escape(sym_escaped)
 
+        # This is to handle interpolated heads -- Setfield treats them differently:
+        # julia> Setfield.parse_obj_lens(@q $name.a)
+        # (:($(Expr(:escape, :_))), :((Setfield.compose)($(Expr(:escape, :name)), (Setfield.PropertyLens){:a}())))
+        # julia> Setfield.parse_obj_lens(@q x.a)
+        # (:($(Expr(:escape, :x))), :((Setfield.compose)((Setfield.PropertyLens){:a}())))
+        if sym != :_
+            sym = QuoteNode(sym)
+        else
+            sym = lens.args[2]
+            lens = Expr(:call, lens.args[1], lens.args[3:end]...)
+        end
+
         if concretize
             return :(
-                $(AbstractPPL.VarName){$(QuoteNode(sym))}(
+                $(AbstractPPL.VarName){$sym}(
                     $(AbstractPPL.concretize)($lens, $sym_escaped)
                 )
             )
         elseif Setfield.need_dynamic_lens(expr)
             error("Variable name `$(expr)` is dynamic and requires concretization!")
         else
-            :($(AbstractPPL.VarName){$(QuoteNode(sym))}($lens))
+            :($(AbstractPPL.VarName){$sym}($lens))
         end
+    elseif Meta.isexpr(expr, :$, 1)
+        return :($(AbstractPPL.VarName){$(esc(expr.args[1]))}())
     else
         error("Malformed variable name `$(expr)`!")
     end