@rand: fix missing esc

Author: rfourquet

src/macros.jl

@@ -3,47 +3,53 @@ macro rand(exp)
 end
 
 function rand_macro(ex)
-    ex isa Expr &&
-        ex.head ∈ (:(=), :function) &&
-        ex.args[1].head == :call &&
-        ex.args[1].args[1] == :rand || throw(ArgumentError(
+    ex isa Expr && ex.head ∈ (:(=), :function) ||
+        throw(ArgumentError("@rand requires an expression defining `rand`"))
+    sig = ex.args[1]
+    body = ex.args[2]
+
+    sig.head == :call &&
+        sig.args[1] == :rand || throw(ArgumentError(
             "@rand requires a function expression defining `rand`"))
 
-    name = ex.args[1].args[2].args[1] # x
-    namefull = ex.args[1].args[2] # x::X
-    @assert namefull.head == :(::)
-    namefull.args[2] = esc(namefull.args[2])
+    argname = sig.args[2].args[1] # x
+    namefull = sig.args[2] # x::X
+    @assert namefull.head == :(::) # TODO: throw exception
 
     sps = Any[] # sub-samplers
-    body = samplerize!(sps, ex.args[2], name)
+    rng = gensym()
+    body = samplerize!(sps, body, argname, rng)
     istrivial = isempty(sps)
     # rand -> Base.rand
-    ex = Expr(ex.head, Expr(:call, :(Random.rand),
-                            :(rng::AbstractRNG),
-                            as_sampler(ex.args[1].args[2], istrivial)),
-              body)
+
+    exsig = Expr(:call,
+                 :(Random.rand),
+                 :($(esc(rng))::AbstractRNG),
+                 esc(as_sampler(namefull, istrivial)))
+
+    ex = Expr(ex.head, exsig, esc(body))
 
     sp = if istrivial
         # we explicitly define Sampler even in the trivial case to handle
         # redefinitions, where the old rand/sampler pair (for SamplerSimple)
         # is overwritten by a new one (for SamplerTrivial)
         quote
             Random.Sampler(::Type{RNG}, n::Repetition) where {RNG<:AbstractRNG} =
-                SamplerTrivial($name)
+                SamplerTrivial($(esc(argname)))
         end
     else
         quote
             Random.Sampler(::Type{RNG}, n::Repetition) where {RNG<:AbstractRNG} =
-                SamplerSimple($name, tuple(SP))
+                SamplerSimple($(esc(argname)), tuple(SP))
         end
     end
 
     # insert x::X in the argument list, between RNG and n::Repetition
-    insert!(sp.args[2].args[1].args[1].args, 3, namefull)
+    insert!(sp.args[2].args[1].args[1].args, 3, esc(namefull))
 
     # insert inner samplers
     if !istrivial
-        SP = [Expr(:call, :Sampler, :RNG, x, :n) for x in sps]
+        SP = [Expr(:call, :Sampler, :RNG, esc(x), :n) for x in sps]
         @assert :SP == pop!(sp.args[2].args[2].args[2].args[3].args)
         append!(sp.args[2].args[2].args[2].args[3].args, SP)
     end
@@ -55,14 +61,14 @@ function rand_macro(ex)
 end
 
 function as_sampler(ex, istrivial)
-    t = istrivial ? :SamplerTrivial : :SamplerSimple
+    t = istrivial ? :(RandomExtensions.SamplerTrivial) : :(RandomExtensions.SamplerSimple)
     Expr(:(::),
          ex.args[1],
          Expr(:curly, t,
               Expr(:(<:), ex.args[2])))
 end
 
-function samplerize!(sps, ex, name)
+function samplerize!(sps, ex, name, rng)
     if ex == name
         # not within a rand() call
         return Expr(:ref, name) # name -> name[]
@@ -72,8 +78,8 @@ function samplerize!(sps, ex, name)
         # TODO: handle Repetition == Val(Inf) for arrays
         push!(sps, ex.args[2])
         i = length(sps)
-        Expr(:call, :rand, :rng, :($name.data[$i]), ex.args[3:end]...)
+        Expr(:call, :rand, rng, :($name.data[$i]), ex.args[3:end]...)
     else
-        Expr(ex.head, map(e -> samplerize!(sps, e, name), ex.args)...)
+        Expr(ex.head, map(e -> samplerize!(sps, e, name, rng), ex.args)...)
     end
 end

test/runtests.jl

@@ -574,9 +574,10 @@ end
 Base.eltype(::Type{Die}) = Int
 
 @testset "@rand" begin
-    d = Die(6)
-    rng = MersenneTwister()
+    # rng0 to be sure rng is not accessed in `esc`aped body of @rand
+    rng0 = MersenneTwister()
 
+    d = Die(6)
     @rand function rand(d::Die)
         7
     end
@@ -585,8 +586,8 @@ Base.eltype(::Type{Die}) = Int
     # redefinition
     @rand rand(d::Die) = rand(1:d.n)
     @test rand(d) ∈ 1:6
-    @test rand(rng, d) ∈ 1:6
-    @test all(∈(1:6), rand(rng, d, 10))
+    @test rand(rng0, d) ∈ 1:6
+    @test all(∈(1:6), rand(rng0, d, 10))
     @test eltype(rand(d, 3)) == Int
 
     # redefinition (multiple inner samplers)
@@ -601,4 +602,9 @@ Base.eltype(::Type{Die}) = Int
     # redefinition back to SamplerTrivial
     @rand rand(d::Die) = d.n
     @test all(==(6), rand(d, 100))
+
+    # test esc-correctness
+    VAR = 100
+    @rand rand(d::Die) = rand(VAR+1:VAR+d.n) - VAR
+    @test all(∈(1:6), rand(d, 100))
 end