Use @eval macro, add support for tracing arbitrary Functions.

Author: ztangent

README.md

@@ -1,7 +1,8 @@
 # TracedRandom.jl
 
 Allows for the optional specification of traced addresses (i.e. variable names)
-in calls to `rand` and other primitive functions in `Random`:
+in calls to `rand` and other primitive functions in `Random`. Providing this
+information allows ordinary Julia code to be "probabilistic-programming-ready".
 
 ```julia
 julia> rand(:u, Float64, 10)
@@ -21,9 +22,17 @@ julia> randperm(:perm)
  4
 ```
 
-By default, the addresses (`:x`, `:z` and `:perm` in the examples above)
-are ignored, but they can be intercepted via meta-programming by
-probabilistic programming systems such as [`Gen`](https://www.gen.dev/) and
-[`Jaynes`](https://femtomc.github.io/Jaynes.jl/) in order to perform inference
-over the addressed random variables. Addresses can be specified as `Symbol`s,
-or as pairs from symbols to other types (`Pair{Symbol,<:Any}`).
+In addition, a call to some `fn::Function` can be annotated with an address
+by wrapping it in `rand` call:
+```julia
+julia> gaussian_mixture(μs) = randn(:z) + μs[rand(:k, 1:length(μs))]
+julia> rand(:x, gaussian_mixture, [1, 10])
+9.594800995267331
+```
+
+By default, the addresses (`:x`, `:z`, `:k` and `:perm` in the examples above)
+are ignored, but they can be intercepted via meta-programming
+(see [`Genify.jl`](https://github.com/probcomp/Genify.jl])) to support inference
+in probabilistic programming systems such as [`Gen`](https://www.gen.dev/).
+Addresses can be specified as `Symbol`s, or as pairs from symbols
+to other types (`Pair{Symbol,<:Any}`).

src/TracedRandom.jl

@@ -8,6 +8,9 @@ export rand!, randn!, randexp, randexp!,
 
 const Address = Union{Symbol,Pair{Symbol}}
 
+Base.rand(::Address, fn::Function, args...) = fn(args...)
+Base.rand(::AbstractRNG, ::Address, fn::Function, args...) = fn(args...)
+
 Base.rand(::Address, args...) = rand(args...)
 Base.rand(rng::AbstractRNG, ::Address, args...) = rand(rng, args...)
 
@@ -18,8 +21,8 @@ for fn in (:rand!, :randn!, :randexp, :randexp!, :bitrand, :randstring,
            :randsubseq, :randsubseq!, :shuffle, :shuffle!,
            :randperm, :randperm!, :randcycle, :randcycle!)
     fn = GlobalRef(Random, fn)
-    eval(:($fn(::Address, args...) = $fn(args...)))
-    eval(:($fn(rng::AbstractRNG, ::Address, args...) = $fn(rng, args...)))
+    @eval $fn(::Address, args...) = $fn(args...)
+    @eval $fn(rng::AbstractRNG, ::Address, args...) = $fn(rng, args...)
 end
 
 end

test/runtests.jl

@@ -1,5 +1,16 @@
 using TracedRandom, Random, Test
 
+@testset "Addressed function calls" begin
+
+gaussian_mixture(μs) = randn(:z) + μs[rand(:k, 1:length(μs))]
+Random.seed!(0)
+traced = rand(:y, gaussian_mixture, [1, 10, 100])
+Random.seed!(0)
+untraced = gaussian_mixture([1, 10, 100])
+@test traced == untraced
+
+end
+
 @testset "Addressed calls with global RNG" begin
 
 for fn in (rand, randn, randexp, bitrand, randstring)