@rand: use Val(Inf) for sub-samplers generating arrays

Author: rfourquet

src/macros.jl

@@ -30,11 +30,12 @@ function rand_macro(ex)
     namefull = sig.args[2] # x::X
     @assert namefull.head == :(::) # TODO: throw exception
 
-    sps = Any[] # sub-samplers
+    # sub-samplers; second argument true forces Val(Inf) for the sampler
+    sps = Pair{<:Any,Bool}[]
+
     rng = gensym()
     body = samplerize!(sps, body, argname, rng)
     istrivial = isempty(sps)
-    # rand -> Base.rand
 
     exsig = Expr(:call,
                  :(Random.rand),
@@ -69,7 +70,7 @@ function rand_macro(ex)
 
     # insert inner samplers
     if !istrivial
-        SP = [Expr(:call, :Sampler, :RNG, esc(x), :n) for x in sps]
+        SP = [Expr(:call, :Sampler, :RNG, esc(x), many ? Val{Inf}() : :n) for (x, many) 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
@@ -95,8 +96,9 @@ function samplerize!(sps, ex, name, rng)
     end
     ex isa Expr || return ex
     if ex.head == :call && ex.args[1] == :rand
-        # TODO: handle Repetition == Val(Inf) for arrays
-        push!(sps, ex.args[2])
+        # we assume that if rand has more than one arg, we want
+        # a Val(Inf) sampler (e.g. rand(1:9, 2, 3)
+        push!(sps, ex.args[2] => length(ex.args) > 2)
         i = length(sps)
         Expr(:call, :rand, rng, :($name.data[$i]), ex.args[3:end]...)
     else

test/runtests.jl

@@ -635,6 +635,15 @@ Base.eltype(::Type{DieT{T}}) where {T} = T
     @rand rand(d::Die) = d.n
     @test all(==(6), rand(d, 100))
 
+    # redefinition (Val(Inf) iff rand calls with 2+ arguments)
+    @rand rand(d::Die) = (rand("asd") + rand("asd", 3)[1]; 1)
+    s = Sampler(MersenneTwister, d, Val(1))
+    @test s.data[1] isa Random.SamplerSimple{String}
+    @test s.data[2] isa Random.SamplerSimple{Vector{Char}} # "proof" of Val(Inf)
+    s = Sampler(MersenneTwister, d, Val(Inf))
+    @test s.data[1] isa Random.SamplerSimple{Vector{Char}}
+    @test s.data[2] isa Random.SamplerSimple{Vector{Char}}
+
     # test esc-correctness
     VAR = 100
     @rand rand(d::Die) = rand(VAR+1:VAR+d.n) - VAR