Implement nested indexing

Author: wouterwln

Project.toml

@@ -33,7 +33,7 @@ Distributions = "0.25"
 Documenter = "1.0"
 GraphPlot = "0.5, 0.6"
 MacroTools = "0.5"
-MetaGraphsNext = "0.6, 0.7"
+MetaGraphsNext = "~0.7.0"
 NamedTupleTools = "0.14"
 Static = "0.8, 1"
 StaticArrays = "1.6"

src/model_macro.jl

@@ -465,19 +465,30 @@ function proxy_args_lhs_eq_rhs(lhs, rhs)
     return :($lhs = $(proxy_args_rhs(rhs)))
 end
 
-function proxy_args_rhs(rhs)
+function recursive_rhs_indexing(rhs)
+    name = rhs
+    while @capture(name, rlabel_[index__] | new(rlabel_[index__]))
+        name = rlabel
+    end
     if isa(rhs, Symbol)
-        return :(GraphPPL.proxylabel($(QuoteNode(rhs)), $rhs, nothing, GraphPPL.False()))
+        return rhs
     elseif @capture(rhs, rlabel_[index__])
-        return :(GraphPPL.proxylabel($(QuoteNode(rlabel)), $rlabel, $(Expr(:tuple, index...)), GraphPPL.False()))
+        return :(GraphPPL.proxylabel($(QuoteNode(name)), $(recursive_rhs_indexing(rlabel)), $(Expr(:tuple, index...)), GraphPPL.False()))
     elseif @capture(rhs, new(rlabel_[index__]))
-        return :(GraphPPL.proxylabel($(QuoteNode(rlabel)), $rlabel, $(Expr(:tuple, index...)), GraphPPL.True()))
+        return :(GraphPPL.proxylabel($(QuoteNode(name)), $(recursive_rhs_indexing(rlabel)), $(Expr(:tuple, index...)), GraphPPL.True()))
     elseif @capture(rhs, rlabel_...)
         return :(GraphPPL.proxylabel($(QuoteNode(rlabel)), GraphPPL.Splat($rlabel), nothing, GraphPPL.False())...)
     end
     return :(GraphPPL.proxylabel(:anonymous, $rhs, nothing, GraphPPL.False()))
 end
 
+function proxy_args_rhs(rhs)
+    if isa(rhs, Symbol)
+        return :(GraphPPL.proxylabel($(QuoteNode(rhs)), $rhs, nothing, GraphPPL.False()))
+    end
+    return recursive_rhs_indexing(rhs)
+end
+
 """
     combine_args(args::Vector, kwargs::Nothing)
 

test/graph_construction_tests.jl

@@ -1984,3 +1984,22 @@ end
     @test GraphPPL.getname(model[normal_node, x[2]]) == :σ
     @test GraphPPL.getname(model[normal_node, y]) == :out
 end
+
+@testitem "Multiple indices in rhs statement" begin
+    using Distributions
+    using GraphPPL
+    import GraphPPL: @model, create_model, datalabel, NodeCreationOptions, neighbors
+
+    @model function multiple_indices(prior_params, y)
+        x ~ Normal(prior_params[1][1], prior_params[1][2])
+        y ~ Normal(x, 1.0)
+    end
+    model = create_model(multiple_indices(prior_params = [[1, 2]])) do model, ctx
+        y = datalabel(model, ctx, NodeCreationOptions(kind = :data), :y, rand())
+        return (y = y,)
+    end
+
+    @test length(collect(filter(as_node(Normal), model))) == 2
+    @test length(collect(filter(as_variable(:y), model))) == 1
+    @test length(collect(filter(as_variable(:x), model))) == 1
+end

test/model_macro_tests.jl

@@ -1181,6 +1181,43 @@ end
     @test_expression_generating apply_pipeline(input, convert_anonymous_variables) output
 end
 
+@testitem "proxy_args_rhs" begin
+    import GraphPPL: proxy_args_rhs, apply_pipeline, recursive_rhs_indexing
+
+    include("testutils.jl")
+
+    # Test 1: Input expression with a function call in rhs arguments
+    input = :x
+    output = quote
+        GraphPPL.proxylabel(:x, x, nothing, GraphPPL.False())
+    end
+    @test_expression_generating proxy_args_rhs(input) output
+
+    input = quote
+        x[1]
+    end
+    output = quote
+        GraphPPL.proxylabel(:x, x, (1,), GraphPPL.False())
+    end
+    @test_expression_generating proxy_args_rhs(input) output
+
+    input = quote
+        x[1, 2]
+    end
+    output = quote
+        GraphPPL.proxylabel(:x, x, (1, 2), GraphPPL.False())
+    end
+    @test_expression_generating proxy_args_rhs(input) output
+
+    input = quote
+        x[1][1]
+    end
+    output = quote
+        GraphPPL.proxylabel(:x, GraphPPL.proxylabel(:x, x, (1,), GraphPPL.False()), (1,), GraphPPL.False())
+    end
+    @test_expression_generating proxy_args_rhs(input) output
+end
+
 @testitem "add_get_or_create_expression" begin
     import GraphPPL: add_get_or_create_expression, apply_pipeline
 
@@ -2048,7 +2085,7 @@ end
     @test !isnothing(GraphPPL.create_model(somemodel(a = 1, b = 2)))
 end
 
-@testitem "model should warn users against incorrect usages of `=` operator with random variables" begin 
+@testitem "model should warn users against incorrect usages of `=` operator with random variables" begin
     using GraphPPL, Distributions
     import GraphPPL: @model
 
@@ -2058,5 +2095,7 @@ end
         y ~ Normal(0, t)
     end
 
-    @test_throws "One of the arguments to `exp` is of type `GraphPPL.VariableRef`. Did you mean to create a new random variable with `:=` operator instead?" GraphPPL.create_model(somemodel())
+    @test_throws "One of the arguments to `exp` is of type `GraphPPL.VariableRef`. Did you mean to create a new random variable with `:=` operator instead?" GraphPPL.create_model(
+        somemodel()
+    )
 end