Create call trace graphs for julia calls. Based mostly on code and input by Tim Holy (https://github.com/cstjean/TraceCalls.jl/issues/61)

TraceCalls2.jl

module TraceCalls2
using JuliaInterpreter, OrderedCollections, LightGraphs

const callchains = OrderedSet{Vector{Method}}()
const modules = Set{Module}()

function callchain(frame::JuliaInterpreter.Frame)
   chain = Method[]
   sc = JuliaInterpreter.scopeof(frame)
   while sc isa Method
       push!(chain, sc)
       frame = frame.caller
       frame === nothing && break
       sc = JuliaInterpreter.scopeof(frame)
   end
   return chain
end

function log_far!(@nospecialize(recurse), frame, istoplevel::Bool=false)
   chain = callchain(frame)
   chain[1].module ∈ modules && push!(callchains, chain)
   return JuliaInterpreter.finish_and_return!(recurse, frame, istoplevel)
end

function encode_vertices(callchains)
   i = 0
   vertices = Dict{Array{Method}, Int}()
   for chain in callchains
      for ind in length(chain):-1:1
         vert = chain[ind:end]
         haskey(vertices, vert) || (vertices[vert] = (i += 1))
      end
   end
   vertices
end

function getnames(vertices)
   names = Vector{String}(undef, length(vertices))
   for (k,v) in vertices
      names[v] = "$(k[1].module).$(k[1].name)"
   end
   names
end

function construct_graph(callchains)
   vertices = encode_vertices(callchains)
   g = SimpleDiGraph(length(vertices))
   i = 0
   for chain in callchains
      for ind in length(chain)-1:-1:1
         add_edge!(g, vertices[chain[ind+1:end]], vertices[chain[ind:end]])
      end
   end
   g, vertices
end

function tracecall(mods::Tuple, call, arg)
   empty!(callchains)
   empty!(modules)
   for m in mods
      push!(modules, m)
   end
   frame = JuliaInterpreter.enter_call(call, arg);
   log_far!(log_far!, frame, false)
   construct_graph(callchains) #, callchains
end

export tracecall, getnames
end

I just wanted to leave a note here that at some point you and I talked on slack about how one could also do this with Cassette, but the problem was we weren't sure how to get the module that a function is defined in from the function object itself.

I just wanted to note that I found out how you can indeed do that! :) It looks like there's a module pointer on the function's method table object, which seems to be filled out. For example:

julia> module A
       foo() = 2
       end
WARNING: replacing module A.
Main.A

julia> methods(A.foo).mt.module
Main.A

And if you only have the function's type, you can get it by walking back through the instance pointer:

julia> ft = typeof(A.foo)
typeof(Main.A.foo)

julia> methods(ft.instance).mt.module
Main.A

Hope that's useful for you for whatever reason in the future! :)

Thanks, that sounds a lot faster!

I'm just in general having issues getting the information out of the Cassette context metadata. See this example

# this is - unedited - the tracer from the docs
using Cassette

Cassette.@context TraceCtx

function Cassette.overdub(ctx::TraceCtx, args...)
    subtrace = Any[]
    push!(ctx.metadata, args => subtrace)
    if Cassette.canrecurse(ctx, args...)
        newctx = Cassette.similarcontext(ctx, metadata = subtrace)
        return Cassette.recurse(newctx, args...)
    else
        return Cassette.fallback(ctx, args...)
    end
end

trace = Any[]

Running it on Plots.scatter

Cassette.overdub(TraceCtx(metadata = trace), scatter, 1:10)

gives

 (Plots.scatter, 1:10) => Any[
    (NamedTuple,) => Any[
        (Core.apply_type, Tuple) => Any[],
        (Core.apply_type, NamedTuple, (), Tuple{}) => Any[],
        (NamedTuple{(),Tuple{}}, ()) => Any[
            (Core.apply_type, NamedTuple, (), Tuple{}) => Any[]
        ]
     ],
     (pairs, NamedTuple()) => Any[
        (keys, NamedTuple()) => Any[],
        (Base.Iterators.Pairs, NamedTuple(), ()) => Any[
            (Core.apply_type, Base.Iterators.Pairs, Union{}, Union{}, Tuple{}, NamedTuple{(),Tuple{}}) => Any[],
            (fieldtype, Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}, 1) => Any[],
            (convert, NamedTuple{(),Tuple{}}, NamedTuple()) => Any[],
            (fieldtype, Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}, 2) => Any[]
        ]
      ],
      (tuple, Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}(), Plots.scatter) => Any[]
  ]

Almost no actual information on the call trace.