Can Julia macros be used to generate code based on specific function implementation?

Question

I am fairly new to Julia and I am learning about metaprogramming.

I would like to write a macro that receive in input a function and returns another function based on the implementation details of its input.

For example given:

function f(x)
    x + 100
end

function g(x)
    f(x)*x
end

function h(x)
    g(x)-0.5*f(x)
end

I would like to write a macro that returns something like that:

function h_traced(x) 
   f = x + 100
   println("loc 1 x: ", x)
   g = f * x
   println("loc 2 x: ", x)
   res = g - 0.5 * f
   println("loc 3 x: ", x)

Now both code_lowered and code_typed seems to give me back the AST in the form of CodeInfo, however when I try to use it programmatically in my macro I get empty object.

macro myExpand(f)
    body = code_lowered(f)
    println("myExpand Body lenght: ",length(body))
end

called like this

@myExpand :(h)

however the same call outside the macro works ok.

code_lowered(h)

At last even the following return an empty CodeInfo.

macro myExpand(f)
    body = code_lowered(Symbol("h"))
    println("myExpand Body lenght: ",length(body))
end

This might be incredible trivial but I could not work out myseld why the h symbol does not resolve to the function defined. Am I missing something about the scope of symbols?

Answer 1

You don't need a macro, you need a generated function. They can not only return code (Expr), but also IR (lowered code). Usually, for this kind of thing, people use Base.uncompressed_ast, not code_lowered. Both Cassette and IRTools simplify the implementation for you, in different ways.

The basic idea is:

1. Have a generated function that takes a function and its arguments
2. In that function, get the IR of that function, and modify it to your purposes
3. Return the new IR from the generated function. This will then be compiled and called on the original arguments.

A short demonstration with IRTools:

julia> IRTools.@dynamo function traced(args...)
           ir = IRTools.IR(args...)
           p = IRTools.Pipe(ir)
           for (v, stmt) in p
               IRTools.insertafter!(p, v, IRTools.xcall(println, "loc $v"))
           end
           return IRTools.finish(p)
       end

julia> function h(x)
           sin(x)-0.5*cos(x)
       end
h (generic function with 1 method)

julia> @code_ir traced(h, 1)
1: (%1, %2)
  %3 = Base.getfield(%2, 1)
  %4 = Base.getfield(%2, 2)
  %5 = Main.sin(%4)
  %6 = (println)("loc %3")
  %7 = Main.cos(%4)
  %8 = (println)("loc %4")
  %9 = 0.5 * %7
  %10 = (println)("loc %5")
  %11 = %5 - %9
  %12 = (println)("loc %6")
  return %11

julia> traced(h, 1)
loc %3
loc %4
loc %5
loc %6
0.5713198318738266

The rest is left as an exercise. The numbers of the variables are off, because they are, of course, shifted during the transformation. You'd have to add some bookkeeping for that, or use the substitute function on Pipe in some way (but I never quite understood it). If you need the name of the variables, you can get the IR with slots preserved by using a different method of the IR constructor.

(And now the advertisement: I have written something like this. It's currently quite inefficient, but you might get some ideas from it.)

Answer 2

I find it useful to think about macros as a way to transform an input syntax into an output syntax.

So you could very well define a macro @my_macro such that

@my_macro function h(x)
  g(x)-0.5*f(x)
end

would expand to something like

function h_traced(x)
  println("entering function: x=", x)
  g(x)-0.5*f(x)
end

But to such a macro, h is merely a name, an identifier (technically, a Symbol) that can be transformed into h_traced. h is not the function that is bound to this name (in the same way as x = 2 involves binding a name x, to an integer value 2, but x is not 2; x is merely a name that can be used to refer to 2). In contrast to this, when you call code_lowered(h), h gets evaluated first, and code_lowered is passed its value (which is a function) as argument.

Back to our macro: expanding to an expression that involves the definition of g and f goes way further than mere syntax transformations: we're leaving the purely syntactic domain, since such a transformation would need to "understand" that these are functions, look up their definitions and so on.

You are right to think about code_lowered and friends: this is IMO the adequate level of abstraction for what you're trying to achieve. You should probably look into tools like Cassette.jl or IRTools.jl. That being said, if you're still relatively new to Julia, you might want to get a bit more used to the language before delving too deeply into such topics.