Fix dot macro and avoid eval

Author: devmotion

src/compiler.jl

@@ -284,40 +284,40 @@ function replace_sampler!(model_info)
     return model_info
 end
 
-# The next function is defined that way because .~ gives a parsing error in Julia 1.0
 """
-\"""
     replace_tilde!(model_info)
 
-Replaces `~` expressions with observation or assumption expressions, updating `model_info`.
-\"""
+Replace `~` and `.~` expressions with observation or assumption expressions, updating `model_info`.
+"""
 function replace_tilde!(model_info)
-    ex = model_info[:main_body]
-    ex = MacroTools.postwalk(ex) do x 
-        if @capture(x, @M_ L_ ~ R_) && M == Symbol("@__dot__")
-            generate_dot_tilde(L, R, model_info)
-        else
-            x
-        end
-    end
-    $(VERSION >= v"1.1" ? "ex = MacroTools.postwalk(ex) do x
-        if @capture(x, L_ .~ R_)
-            generate_dot_tilde(L, R, model_info)
-        else
-            x
+    # Apply the `@.` macro first.
+    expr = model_info[:main_body]
+    dottedexpr = MacroTools.postwalk(apply_dotted, expr)
+
+    # Check for tilde operators.
+    tildeexpr = MacroTools.postwalk(dottedexpr) do x
+        # Check dot tilde first.
+        dotargs = getargs_dottilde(x)
+        if dotargs !== nothing
+            L, R = dotargs
+            return generate_dot_tilde(L, R, model_info)
         end
-    end" : "")
-    ex = MacroTools.postwalk(ex) do x
-        if @capture(x, L_ ~ R_)
-            generate_tilde(L, R, model_info)
-        else
-            x
+
+        # Check tilde.
+        args = getargs_tilde(x)
+        if args !== nothing
+            L, R = args
+            return generate_tilde(L, R, model_info)
         end
+
+        return x
     end
-    model_info[:main_body] = ex
+
+    # Update the function body.
+    model_info[:main_body] = tildeexpr
+
     return model_info
 end
-""" |> Meta.parse |> eval
 
 """
     generate_tilde(left, right, model_info)

src/utils.jl

@@ -1,3 +1,53 @@
+"""
+    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)
+
+Return the arguments `L` and `R`, if `x` is an expression of the form `L .~ R` or
+`(~).(L, R)`, or `nothing` otherwise.
+"""
+getargs_dottilde(x) = nothing
+function getargs_dottilde(expr::Expr)
+    # Check if the expression is of the form `L .~ R`.
+    if Meta.isexpr(expr, :call, 3) && expr.args[1] === :.~
+        return expr.args[2], expr.args[3]
+    end
+
+    # Check if the expression is of the form `(~).(L, R)`.
+    if Meta.isexpr(expr, :., 2) && expr.args[1] === :~ &&
+        Meta.isexpr(expr.args[2], :tuple, 2)
+        return expr.args[2].args[1], expr.args[2].args[2]
+    end
+
+    return
+end
+
+"""
+    getargs_tilde(x)
+
+Return the arguments `L` and `R`, if `x` is an expression of the form `L ~ R`, or `nothing`
+otherwise.
+"""
+getargs_tilde(x) = nothing
+function getargs_tilde(expr::Expr)
+    if Meta.isexpr(expr, :call, 3) && expr.args[1] === :~
+        return expr.args[2], expr.args[3]
+    end
+    return
+end
+
 ############################################
 # Julia 1.2 temporary fix - Julia PR 33303 #
 ############################################

test/compiler.jl

@@ -313,14 +313,23 @@ priors = 0 # See "new grammar" test.
         x = randn(100)
         res = sample(vdemo1(x), alg, 250)
 
+        @model vdemo1b(x) = begin
+            s ~ InverseGamma(2,3)
+            m ~ Normal(0, sqrt(s))
+            @. x ~ Normal(m, $(sqrt(s)))
+            return s, m
+        end
+
+        res = sample(vdemo1b(x), alg, 250)
+
         D = 2
         @model vdemo2(x) = begin
             μ ~ MvNormal(zeros(D), ones(D))
-            @. x ~ MvNormal(μ, ones(D))
+            @. x ~ $(MvNormal(μ, ones(D)))
         end
 
         alg = HMC(0.01, 5)
-        res = sample(vdemo2(randn(D,100)), alg, 250)
+        res = sample(vdemo2(randn(D, 100)), alg, 250)
 
         # Vector assumptions
         N = 10