Variables in macro definition

I am developing a macro that lets me write sums in a nice way. A trivial example is:

# some array
x = [1, 2, 3]

# calculate sum
s = @Σ i=1:3 x[i]

I have tried this:

module MyModule

export @Σ

# fancy macro to write sums
macro Σ(limits, expr)
    index = limits.args[1]
    start, stop = limits.args[2].args
    quote
       local out = $(_indexreplace(index, start, expr))
       for i = ($start+1):$stop
           out += $(_indexreplace(index, :i, expr))
       end
       out 
    end
end

# replace the previous index with the loop index inside the macro
_indexreplace(old, new, e::Expr) = Expr(e.head, (_indexreplace(old, new, sub) for sub in e.args)...)
_indexreplace(old, new, e::Symbol) = e == old ? new : e

end

# tests
using MyModule

x = [1, 2, 3]

println(macroexpand(:(@Σ i=1:3 x[i])))

@Σ i=1:3 x[i]

Running this gives me

begin  # /Users/davide/Research/RMCO/julia/LorenzRMCO/test/test_sums.jl, line 10:
    local #2#out = (MyModule.x)[1] # /Users/davide/Research/RMCO/julia/LorenzRMCO/test/test_sums.jl, line 11:
    for #3#i = 1 + 1:3 # /Users/davide/Research/RMCO/julia/LorenzRMCO/test/test_sums.jl, line 12:
        #2#out += (MyModule.x)[#3#i]
    end # /Users/davide/Research/RMCO/julia/LorenzRMCO/test/test_sums.jl, line 14:
    #2#out
end
ERROR: LoadError: UndefVarError: x not defined
 in macro expansion; at /Users/davide/Research/RMCO/julia/LorenzRMCO/test/test_sums.jl:10 [inlined]
 in anonymous at ./<missing>:?
 in include_from_node1(::String) at ./loading.jl:488
 in include_from_node1(::String) at /Users/davide/Software/julia-0.5/usr/lib/julia/sys.dylib:?
 in process_options(::Base.JLOptions) at ./client.jl:262
 in _start() at ./client.jl:318
 in _start() at /Users/davide/Software/julia-0.5/usr/lib/julia/sys.dylib:?
while loading /Users/davide/Research/RMCO/julia/LorenzRMCO/test/test_sums.jl, in expression starting on line 28

which results in an error because the macro code is calling MyModule.x which is not what I want. I have tried escaping the expression with esc, as discussed in the manual, but I can’t seem to find a solution to this. Any help is appreciated.

Thanks,

Davide

In general, you must escape all the user input once and exactly once.

Example

julia> macro Σ(limits::Expr, expr)
           @assert limits.head === :(=)
           @assert length(limits.args) == 2
           index = limits.args[1]
           @assert Meta.isexpr(limits.args[2], :(:))
           start, stop = limits.args[2].args
           quote
               local _start = $(esc(start))
               local _stop = $(esc(stop))
               $(esc(index)) = _start
               local out = $(esc(expr))
               for $(esc(index)) in (_start + 1):_stop
                   out += $(esc(expr))
               end
               out
           end
       end
@Σ (macro with 1 method)

julia> f(x) = @Σ i=1:3 x[i]
f (generic function with 1 method)

julia> f([1, 2, 3])
6
3 Likes

What about
sum(x[i] for i=1:3)

I also tried the generator approach, but it can be much slower than expanding the body directly into a for loop. See this:

module MyModule

export @Σ

# fancy macro to write sums
macro Σ(limits, expr)
    @assert limits.head === :(=)
    @assert length(limits.args) == 2
    index = limits.args[1]
    start, stop = limits.args[2].args
    quote
       local _start = $(esc(start))
       local _stop  = $(esc(stop))
       $(esc(index)) = _start
       local out = $(esc(expr))
       for $(esc(index)) = (_start+1):_stop
           out += $(esc(expr))
       end
       out 
    end
end

end

using MyModule
using BenchmarkTools

N = 400
x = randn(N, N)

fun1(x, j, N) = @Σ i=1:N x[i, j]
fun2(x, j, N) = sum(x[i, j] for i=1:N)

println(@benchmark fun1($x, 2, N))
println(@benchmark fun2($x, 2, N))

On my machine I get

Trial(320.763 ns)
Trial(5.931 μs)

that is twenty times faster.