digbyterrell
digbyterrell

Reputation: 3619

Metaprogramming in Julia -- Splice integer into variable name

Short Question:

I have a Integer variable N and would like to write a macro to produce a single dummy variable i_($N).

An attempt:

@generated function testfunc(N)
     :(i_($N))
end
testfunc(5) # Desired behavior i_5

ERROR: UndefVarError: i_ not defined
 in testfunc at none:2

Longer explanation:

I recently discovered Base.Cartesian in Julia. It has some handy tricks for generating dummy variables for indexing multidimensional arrays.

The @ntuple macro in Cartesian can produce sequences starting at 1. For example, @ntuple 5 k->i_k produces (i_1,i_2,i_3,i_4,i_5). Inside an @generated function, if W=5 then @ntuple ($W) k->i_k will produce the same sequence. This doesn't work: @ntuple 1 k->i_(k+$W)

I can't figure out a way to just produce, e.g. i_3 if N=3 (this could be inside an @generated function).

My end goal is to use @nloops in Cartesian to loop over a sequence of dummy variables and store the results in some storage vector indexed by one of the dummy variables, e.g. something like:

@nloops ($W) i A begin
    # Example generated code for N=2, W=3:
    #    storage[i_2] *= A[i_1, i_2, i_3]
    storage[i_($N)] *= A[(@ntuple ($W) k->i_k)...]
end

Upvotes: 2

Views: 473

Answers (3)

mbauman
mbauman

Reputation: 31342

As the other answers suggest, the key is to explicitly construct the symbol you want and splice it into the generated expression. Also note that you can simplify your A[i_1, i_2, …] expression with @nref.

@nloops ($W) i A begin
    storage[$(symbol(:i_, N)] *= @nref $W A i
end

In my experience, the best way to learn how to use these macros is through the use of macroexpand:

julia> using Base.Cartesian
       W = 3
       N = 2
       macroexpand(:(@nloops ($W) i A begin
                         storage[$(symbol(:i_, N))] *= @nref $W A i
                     end))
quote  # cartesian.jl, line 31:
    for i_3 = 1:size(A,3) # cartesian.jl, line 32:
        nothing # cartesian.jl, line 33:
        begin  # cartesian.jl, line 31:
            for i_2 = 1:size(A,2) # cartesian.jl, line 32:
                nothing # cartesian.jl, line 33:
                begin  # cartesian.jl, line 31:
                    for i_1 = 1:size(A,1) # cartesian.jl, line 32:
                        nothing # cartesian.jl, line 33:
                        begin  # none, line 5:
                            storage[i_2] *= A[i_1,i_2,i_3]
                        end # cartesian.jl, line 34:
                        nothing
                    end
                end # cartesian.jl, line 34:
                nothing
            end
        end # cartesian.jl, line 34:
        nothing
    end
end

Upvotes: 3

ShaoWei Teo
ShaoWei Teo

Reputation: 421

I have a Integer variable N and would like to write a macro to produce a single dummy variable i_($N).

To answer your question directly, you can do this by:-

i_5 = 7

macro testfunc(expr)
    return symbol(:i_, expr)
end

@testfunc 5    #This give us back the value of `i_5`, which is 7
@testfunc(5)   #Another way to write the same macro.

So if we have a function:-

function foo()
    i_3 = 9
    return @testfunc(3)
end

foo() will give us the correct answer of 9.

That said, it may still not be sufficient for what you want, which is to put the macro to use inside a function. In gist, when @testfunc(N) where N = 3 generates i_3, we want to refer to i_3 in the function, but if we write:-

function foo(N)
    i_3 = 9
    return @testfunc(N)
end

foo(3) will give us an UndefVarError: i_N not defined. Not sure whether there is any way to rectify it. There are 2 possibilities I can think of to workaround this, one is to write your process as a script in the global scope, another, is to put the whole function in quote.

Upvotes: 2

张实唯
张实唯

Reputation: 2862

I haven't understand what exactly you want, but there are some notes:

  1. @generated arguments are Types in the macro body. eg. in @generated function testfunc(N), N is something like Int64 rather than 3.

  2. if you want to concat a variable name, you may need to construct it by Symbol or something else, but not just quote.

eg.

macro testfunc(N)
    Symbol("i_$N")
end

@testfunc(2) # equivalent to i_2

Upvotes: 3

Related Questions