@rand: handle where clauses in definitions

Author: rfourquet

src/macros.jl

@@ -3,11 +3,17 @@ macro rand(exp)
 end
 
 function rand_macro(ex)
+    whereparams = []
     ex isa Expr && ex.head ∈ (:(=), :function) ||
         throw(ArgumentError("@rand requires an expression defining `rand`"))
     sig = ex.args[1]
     body = ex.args[2]
 
+    if sig.head == :where
+        append!(whereparams, sig.args[2:end])
+        sig = sig.args[1]
+    end
+
     sig.head == :call &&
         sig.args[1] == :rand || throw(ArgumentError(
             "@rand requires a function expression defining `rand`"))
@@ -27,6 +33,10 @@ function rand_macro(ex)
                  :($(esc(rng))::AbstractRNG),
                  esc(as_sampler(namefull, istrivial)))
 
+    if !isempty(whereparams)
+        exsig = Expr(:where, exsig, map(esc, whereparams)...)
+    end
+
     ex = Expr(ex.head, exsig, esc(body))
 
     sp = if istrivial
@@ -43,6 +53,8 @@ function rand_macro(ex)
                 SamplerSimple($(esc(argname)), tuple(SP))
         end
     end
+    @assert sp.args[2].args[1].head == :where
+    append!(sp.args[2].args[1].args, map(esc, whereparams))
 
     # insert x::X in the argument list, between RNG and n::Repetition
     insert!(sp.args[2].args[1].args[1].args, 3, esc(namefull))

test/runtests.jl

@@ -573,6 +573,12 @@ end
 
 Base.eltype(::Type{Die}) = Int
 
+struct DieT{T}
+    n::T
+end
+
+Base.eltype(::Type{DieT{T}}) where {T} = T
+
 @testset "@rand" begin
     # rng0 to be sure rng is not accessed in `esc`aped body of @rand
     rng0 = MersenneTwister()
@@ -607,4 +613,22 @@ Base.eltype(::Type{Die}) = Int
     VAR = 100
     @rand rand(d::Die) = rand(VAR+1:VAR+d.n) - VAR
     @test all(∈(1:6), rand(d, 100))
+
+    # with type parameters
+    d = DieT(6)
+
+    @rand rand(d::DieT{T}) where {T} = 1
+    @test rand(d) == 1
+
+    @rand rand(d::DieT{Int}) = 0
+    @test rand(d) == 0
+
+    d = DieT(0x0)
+    @rand rand(d::DieT{T}) where {T<:UInt8} = rand(typemin(T):typemax(T))
+    @test rand(d) isa UInt8
+
+    d = DieT(true)
+    @rand rand(d::DieT{T}) where {T<:Bool} =
+        T(typemin(T) + rand(typemin(T):typemax(T)))
+    @test rand(d) isa Bool
 end