@rand: accept anonymous functions

Author: rfourquet

src/macros.jl

@@ -4,7 +4,7 @@ end
 
 function rand_macro(ex)
     whereparams = []
-    ex isa Expr && ex.head ∈ (:(=), :function) ||
+    ex isa Expr && ex.head ∈ (:(=), :function, :->) ||
         throw(ArgumentError("@rand requires an expression defining `rand`"))
     sig = ex.args[1]
     body = ex.args[2]
@@ -14,6 +14,14 @@ function rand_macro(ex)
         sig = sig.args[1]
     end
 
+    if ex.head == :function && sig.head == :tuple # anonymous function
+        sig = Expr(:call, :rand, sig.args...)
+    end
+    if ex.head == :->
+        # TODO: check that only one argument is passed
+        sig = Expr(:call, :rand, sig)
+    end
+
     sig.head == :call &&
         sig.args[1] == :rand || throw(ArgumentError(
             "@rand requires a function expression defining `rand`"))
@@ -37,7 +45,7 @@ function rand_macro(ex)
         exsig = Expr(:where, exsig, map(esc, whereparams)...)
     end
 
-    ex = Expr(ex.head, exsig, esc(body))
+    ex = Expr(:function, exsig, esc(body))
 
     sp = if istrivial
         # we explicitly define Sampler even in the trivial case to handle

test/runtests.jl

@@ -588,6 +588,10 @@ Base.eltype(::Type{DieT{T}}) where {T} = T
         7
     end
     @test rand(d) == 7
+    @rand function (d::Die) 7 end
+    @test rand(d) == 7
+    @rand (d::Die) -> 7
+    @test rand(d) == 7
 
     # redefinition
     @rand rand(d::Die) = rand(1:d.n)
@@ -596,11 +600,33 @@ Base.eltype(::Type{DieT{T}}) where {T} = T
     @test all(∈(1:6), rand(rng0, d, 10))
     @test eltype(rand(d, 3)) == Int
 
+    @rand function (d::Die) rand(1:d.n) end
+    @test rand(d) ∈ 1:6
+    @test rand(rng0, d) ∈ 1:6
+    @test all(∈(1:6), rand(rng0, d, 10))
+    @test eltype(rand(d, 3)) == Int
+
+    @rand (d::Die) -> rand(1:d.n)
+    @test rand(d) ∈ 1:6
+    @test rand(rng0, d) ∈ 1:6
+    @test all(∈(1:6), rand(rng0, d, 10))
+    @test eltype(rand(d, 3)) == Int
+
     # redefinition (multiple inner samplers)
     @rand rand(d::Die) = rand(10:10) + rand(1:d.n)
     @test rand(d) ∈ 11:16
     @test all(∈(11:16), rand(d, 10))
 
+    @rand function (d::Die)
+        rand(10:10) + rand(1:d.n)
+    end
+    @test rand(d) ∈ 11:16
+    @test all(∈(11:16), rand(d, 10))
+
+    @rand (d::Die) -> rand(10:10) + rand(1:d.n)
+    @test rand(d) ∈ 11:16
+    @test all(∈(11:16), rand(d, 10))
+
     # redefinition access to argument not within rand call
     @rand rand(d::Die) = rand(1:1) + d.n
     @test all(==(7), rand(d, 10))
@@ -620,6 +646,19 @@ Base.eltype(::Type{DieT{T}}) where {T} = T
     @rand rand(d::DieT{T}) where {T} = 1
     @test rand(d) == 1
 
+    @rand function rand(d::DieT{T}) where {T}
+        2
+    end
+    @test rand(d) == 2
+
+    @rand function (d::DieT{T}) where {T}
+        3
+    end
+    @test rand(d) == 3
+
+    @rand ((d::DieT{T}) where {T}) -> 4
+    @test rand(d) == 4
+
     @rand rand(d::DieT{Int}) = 0
     @test rand(d) == 0