Can Julia's macros access nested functions?

Question

I am playing around with Julia's macros. One thing I am particularly curious about is the ability to extract a function's reachable call graph without having to fully compile the code. By 'reachable call graph', I mean all of the functions found inside the body of a function that are potentially reachable, and the functions in those functions, and so on.

For example:

do_something(x::Int) = println(x*2)
do_something(x::String) = println(x)

function foo(a, b)::Bool
    do_something_with(a)
    do_something_with(b)
    return true
end

# Ideally something like this would be accessible from Base
functions_in_codebase = [:do_something, :foo]

# Not correct, but gets the idea across
macro access_call_graph(f)
    tokens = f.text.tokens
    graph = [f]
    for t in tokens
        go_deeper = t in functions_in_codebase
        !go_deeper && append!(graph, access_call_graph(get_function_for(t))...)
    end
    return graph
end

access_call_graph(foo)
# Should get something like the following, nesting notwithstanding:
# foo
# do_something, do_something

This is pretty janky, but being able to access the call graph at parse-time, even if only in terms of potentially reachable functions, would be extremely useful for what I'm trying to do. If I have to fully compile the code in order for something like this to work, that would largely defeat the benefit.

Is something like this possible?

It won't work the way you wrote it so far. For one, macros aren't functions, their calls look like `@access_call_graph(foo)`. For another, macros only see expressions at parse-time, so that call only sees the Symbol `:foo`, not the function. `:foo` definitely won't have a `.text.tokens` field. The way you wrote `access_call_graph` looks more like a function that operates at run-time. This reminds me of the function `code_warntype`. The macro `@code_warntype` is more well-known, but all that macro does is process an expression to an expression that calls the function version. — BatWannaBe, Jan 03 '22 at 22:36

score 3 · Accepted Answer · answered Jan 03 '22 at 22:42

3

@code_lowered could be useful:

julia> code = @code_lowered foo(1,"2")
CodeInfo(
1 ─ %1 = Main.Bool
│        Main.do_something_with(a)
│        Main.do_something_with(b)
│   %4 = Base.convert(%1, true)
│   %5 = Core.typeassert(%4, %1)
└──      return %5
)

To get function calls you could do:

julia> code.first.code
5-element Vector{Any}:
 :(Main.do_something_with)
 :((%1)(_2))
 :(Main.do_something_with)
 :((%3)(_3))
 :(return true)

answered Jan 03 '22 at 22:42

Przemyslaw Szufel

40,002
3
32
62

3

This, in combination with @BatWannaBe's comment above is essentially the correct answer, but it gets complicated quickly, since you need to recur for every _method_, not function. I'm just gonna point to my package [IRTracker.jl](https://github.com/TuringLang/IRTracker.jl), which isn't maintained anymore, but does pretty much what the OP wants and can serve as an example of how to write something like it using [IRTools.jl](https://github.com/FluxML/IRTools.jl). And I'd recommend the latter as a starting point. – phipsgabler Jan 04 '22 at 09:54
2

@phipsgabler Just expanding on your first sentence, the reason OP needs to recur over each method instead of just the function is because different methods of the same function can call different nested functions. The call graph for `@access_call_graph f(y)` is going to look different depending on y's type: `f(x::Int) = g(x)` vs `f(x::Float64) = h(x)`. So while OP may not want full compilation, OP'll need the type inference part to find nested methods. I suppose a version that works on a *function* `@access_call_graphs f` could iterate over all its existing methods `Base.methods(f)`. – BatWannaBe Jan 04 '22 at 12:15

Can Julia's macros access nested functions?

1 Answers1