Julia: inject code into function - macros

Question

I would like to inject code into a function. For concreteness, consider a simple simulater:
function simulation(A, x)
for t in 1:1000
z = randn(3)
x = A*x + z
end
end
Sometimes I would like to record the values of x every ten time-steps, sometimes the values of z every 20 time-steps, and sometimes I don't want to record any values. I could, of course, put some flags as arguments to the function, and have some if-else statements. But I would like to rather keep the simulation code clean, and only inject a piece of code like
if t%10 == 0
append!(rec_z, z)
end
into particular places of the function whenever I need it. For that, I'd like to write a macro such that monitoring a particular value becomes
#monitor(:z, 10)
simulation(A, x)
Is that possible with Julia's Metaprogramming capabilities?

Answer 1

No, you cannot use metaprogramming to inject code into an already-written function. Metaprogramming can only do things that you could directly write yourself at precisely the location where the macro itself is written. That means that a statement like:
#monitor(:z, 10); simulation(A, x)
cannot even modify the simulation(A, x) function call. It can only expand out to some normal Julia code that runs before simulation is called. You could, perhaps, include the simulation function call as an argument to the macro, e.g., #monitor(:z, 10, simulation(A, x)), but now all the macro can do is change the function call itself. It still cannot "go back" and add new code to a function that was already written.
You could, however, carefully and meticulously craft a macro that takes the function definition body and modifies it to add your debug code, e.g.,
#monitor(:z, 10, function simulation(A, x)
for t in 1:1000
# ...
end
end)
But now you must write code in the macro that traverses the code in the function body, and injects your debug statement at the correct place. This is not an easy task. And it's even harder to write in a robust manner that wouldn't break the moment you modified your actual simulation code.
Traversing code and inserting it is a much easier task for you to do yourself with an editor. A common idiom for debugging statements is to use a one-liner, like this:
const debug = false
function simulation (A, x)
for t in 1:1000
z = rand(3)
x = A*x + z
debug && t%10==0 && append!(rec_z, z)
end
end
What's really cool here is that by marking debug as constant, Julia is able to completely optimize away the debugging code when it's false — it doesn't even appear in the generated code! So there is no overhead when you're not debugging. It does mean, however, that you have to restart Julia (or reload the module it's in) for you to change the debug flag. Even when debug isn't marked as const, I cannot measure any overhead for this simple loop. And chances are, your loop will be more complicated than this one. So don't worry about performance here until you actually double-check that it's having an effect.

Answer 2

You might be interested in this which i just whipped up. It doesn't QUITE do what you are doing, but it's close. Generally safe and consistent places to add code are the beginning and end of code blocks. These macros allow you to inject some code in those location (and even pass code parameters!)
Should be useful for say toggle-able input checking.
#cleaninject.jl
#cleanly injects some code into the AST of a function.
function code_to_inject()
println("this code is injected")
end
function code_to_inject(a,b)
println("injected code handles $a and $b")
end
macro inject_code_prepend(f)
#make sure this macro precedes a function definition.
isa(f, Expr) || error("checkable macro must precede a function definition")
(f.head == :function) || error("checkable macro must precede a function definition")
#be lazy and let the parser do the hard work.
b2 = parse("code_to_inject()")
#inject the generated code into the AST.
unshift!(f.args[2].args, b2)
#return the escaped function to the parser so that it generates the new function.
return Expr(:escape, f)
end
macro inject_code_append(f)
#make sure this macro precedes a function definition.
isa(f, Expr) || error("checkable macro must precede a function definition")
(f.head == :function) || error("checkable macro must precede a function definition")
#be lazy and let the parser do the hard work.
b2 = parse("code_to_inject()")
#inject the generated code into the AST.
push!(f.args[2].args, b2)
#return the escaped function to the parser so that it generates the new function.
return Expr(:escape, f)
end
macro inject_code_with_args(f)
#make sure this macro precedes a function definition.
isa(f, Expr) || error("checkable macro must precede a function definition")
(f.head == :function) || error("checkable macro must precede a function definition")
#be lazy and let the parser do the hard work.
b2 = parse(string("code_to_inject(", join(f.args[1].args[2:end], ","), ")"))
#inject the generated code into the AST.
unshift!(f.args[2].args, b2)
#return the escaped function to the parser so that it generates the new function.
return Expr(:escape, f)
end
################################################################################
# RESULTS
#=
julia> #inject_code_prepend function p()
println("victim function")
end
p (generic function with 1 method)
julia> p()
this code is injected
victim function
julia> #inject_code_append function p()
println("victim function")
end
p (generic function with 1 method)
julia> p()
victim function
this code is injected
julia> #inject_code_with_args function p(a, b)
println("victim called with $a and $b")
end
p (generic function with 2 methods)
julia> p(1, 2)
injected code handles 1 and 2
victim called with 1 and 2
=#