Do not modify interpolated expressions

Author: devmotion

src/compiler.jl

@@ -69,9 +69,17 @@ end
 
 To generate a `Model`, call `model_generator(x_value)`.
 """
-macro model(input_expr)
-    Base.replace_ref_end!(input_expr) |> build_model_info |> replace_tilde! |> replace_vi! |> 
-        replace_logpdf! |> replace_sampler! |> build_output
+macro model(expr)
+    esc(model(expr))
+end
+
+function model(expr)
+    modelinfo = build_model_info(expr)
+
+    # Generate main body
+    modelinfo[:main_body] = generate_mainbody(modelinfo)
+
+    return build_output(modelinfo)
 end
 
 """
@@ -170,62 +178,55 @@ function build_model_info(input_expr)
     return model_info
 end
 
-
 """
-    replace_vi!(model_info)
+    generate_mainbody([expr, ]modelinfo)
 
-Replaces `@varinfo()` expressions with a handle to the `VarInfo` struct.
+Generate the body of the main evaluation function.
 """
-function replace_vi!(model_info)
-    ex = model_info[:main_body]
-    vi = model_info[:main_body_names][:vi]
-    ex = MacroTools.postwalk(ex) do x
-        if @capture(x, @varinfo())
-            vi
-        else
-            x
-        end
-    end
-    model_info[:main_body] = ex
-    return model_info
-end
+generate_mainbody(modelinfo) = generate_mainbody(modelinfo[:main_body], modelinfo)
 
-"""
-    replace_logpdf!(model_info)
+generate_mainbody(x, modelinfo) = x
+function generate_mainbody(expr::Expr, modelinfo)
+    # Do not touch interpolated expressions
+    expr.head === :$ && return expr.args[1]
 
-Replaces `@logpdf()` expressions with the value of the accumulated `logpdf` in the `VarInfo` struct.
-"""
-function replace_logpdf!(model_info)
-    ex = model_info[:main_body]
-    vi = model_info[:main_body_names][:vi]
-    ex = MacroTools.postwalk(ex) do x
-        if @capture(x, @logpdf())
-            :($(vi).logp[])
-        else
-            x
+    # Apply the `@.` macro first.
+    if Meta.isexpr(expr, :macrocall) && length(expr.args) > 1 &&
+        expr.args[1] === Symbol("@__dot__")
+        return generate_mainbody(Base.Broadcast.__dot__(expr.args[end]), modelinfo)
+    end
+
+    # Modify macro calls.
+    if Meta.isexpr(expr, :macrocall) && !isempty(expr.args)
+        name = expr.args[1]
+        if name === Symbol("@varinfo")
+            return modelinfo[:main_body_names][:vi]
+        elseif name === Symbol("@logpdf")
+            return :($(modelinfo[:main_body_names][:vi]).logp[])
+        elseif name === Symbol("@sampler")
+            return :($(modelinfo[:main_body_names][:sampler]))
         end
     end
-    model_info[:main_body] = ex
-    return model_info
-end
 
-"""
-    replace_sampler!(model_info)
+    # Modify dotted tilde operators.
+    args_dottilde = getargs_dottilde(expr)
+    if args_dottilde !== nothing
+        L, R = args_dottilde
+        return Base.remove_linenums!(generate_dot_tilde(generate_mainbody(L, modelinfo),
+                                                        generate_mainbody(R, modelinfo),
+                                                        modelinfo))
+    end
 
-Replaces `@sampler()` expressions with a handle to the sampler struct.
-"""
-function replace_sampler!(model_info)
-    ex = model_info[:main_body]
-    spl = model_info[:main_body_names][:sampler]
-    ex = MacroTools.postwalk(ex) do x
-        if @capture(x, @sampler())
-            spl
-        else
-            x
-        end
+    # Modify tilde operators.
+    args_tilde = getargs_tilde(expr)
+    if args_tilde !== nothing
+        L, R = args_tilde
+        return Base.remove_linenums!(generate_tilde(generate_mainbody(L, modelinfo),
+                                                    generate_mainbody(R, modelinfo),
+                                                    modelinfo))
     end
-    model_info[:main_body] = ex
-    return model_info
+
+    return Expr(expr.head, map(x -> generate_mainbody(x, modelinfo), expr.args)...)
 end
 
 """
@@ -443,7 +444,7 @@ function build_output(model_info)
     generator_kw_form = isempty(args) ? () : (:($generator(;$(args...)) = $generator($(arg_syms...))),)
     model_gen_constructor = :($(DynamicPPL.ModelGen){$(Tuple(arg_syms))}($generator, $defaults_nt))
 
-    ex = quote
+    return quote
         function $evaluator(
             $model::$(DynamicPPL.Model),
             $vi::$(DynamicPPL.VarInfo),
@@ -459,8 +460,6 @@ function build_output(model_info)
 
         $model_gen = $model_gen_constructor
     end
-
-    return esc(ex)
 end
 
 

src/utils.jl

@@ -1,17 +1,3 @@
-"""
-    apply_dotted(x)
-
-Apply the transformation of the `@.` macro if `x` is an expression of the form `@. X`.
-"""
-apply_dotted(x) = x
-function apply_dotted(expr::Expr)
-    if Meta.isexpr(expr, :macrocall) && length(expr.args) > 1 &&
-        expr.args[1] === Symbol("@__dot__")
-        return Base.Broadcast.__dot__(expr.args[end])
-    end
-    return expr
-end
-
 """
     getargs_dottilde(x)
 

src/varname.jl

@@ -187,8 +187,10 @@ end
 vinds(expr::Symbol) = Expr(:tuple)
 function vinds(expr::Expr)
     if Meta.isexpr(expr, :ref)
-        last = Expr(:tuple, expr.args[2:end]...)
-        init = vinds(expr.args[1]).args
+        ex = copy(expr)
+        Base.replace_ref_end!(ex)
+        last = Expr(:tuple, ex.args[2:end]...)
+        init = vinds(ex.args[1]).args
         return Expr(:tuple, init..., last)
     else
         throw("VarName: Mis-formed variable name $(expr)!")

test/compiler.jl

@@ -8,6 +8,14 @@ Random.seed!(129)
 
 priors = 0 # See "new grammar" test.
 
+macro custom(expr)
+    (Meta.isexpr(expr, :call, 3) && expr.args[1] === :~) ||
+        error("incorrect macro usage")
+    quote
+        $(esc(expr.args[2])) = 0.0
+    end
+end
+
 @testset "compiler.jl" begin
     @testset "assume" begin
         @model test_assume() = begin
@@ -550,4 +558,12 @@ priors = 0 # See "new grammar" test.
         @test haskey(vi3.metadata, :y)
         @test vi3.metadata.y.vns[1] == VarName(:y, ((1,),))
     end
+    @testset "custom tilde" begin
+        @model demo() = begin
+            $(@custom m ~ Normal())
+            return m
+        end
+        model = demo()
+        @test all(iszero(model()) for _ in 1:1000)
+    end
 end

test/utils.jl

@@ -1,24 +1,8 @@
 using DynamicPPL
-using DynamicPPL: apply_dotted, getargs_dottilde, getargs_tilde
+using DynamicPPL: getargs_dottilde, getargs_tilde
 
 using Test
 
-@testset "apply_dotted" begin
-    # Some things that are not expressions.
-    @test apply_dotted(:x) === :x
-    @test apply_dotted(1.0) === 1.0
-    @test apply_dotted([1.0, 2.0, 4.0]) == [1.0, 2.0, 4.0]
-
-    # Some expressions.
-    @test apply_dotted(:(x ~ Normal(μ, σ))) == :(x ~ Normal(μ, σ))
-    @test apply_dotted(:((.~)(x, Normal(μ, σ)))) == :((.~)(x, Normal(μ, σ)))
-    @test apply_dotted(:((~).(x, Normal(μ, σ)))) == :((~).(x, Normal(μ, σ)))
-    @test apply_dotted(:(@. x ~ Normal(μ, σ))) == :((~).(x, Normal.(μ, σ)))
-    @test apply_dotted(:(@. x ~ Normal(μ, $(Expr(:$, :(sqrt(v))))))) ==
-        :((~).(x, Normal.(μ, sqrt(v))))
-    @test apply_dotted(:(@~ Normal.(μ, σ))) == :(@~ Normal.(μ, σ))
-end
-
 @testset "getargs_dottilde" begin
     # Some things that are not expressions.
     @test getargs_dottilde(:x) === nothing