Method tracing in Julia?

One of my favourite (and most missed) tools for exploring and debugging code is the trace facility often seen in Lisp. I.e., in Common Lisp, I can just trace any function and get a view of its stack trace:

CL-USER> (trace fib)
(FIB)

CL-USER> (fib 3)
  0: (FIB 3)
    1: (FIB 2)
      2: (FIB 1)
      2: FIB returned 1
      2: (FIB 0)
      2: FIB returned 0
    1: FIB returned 1
    1: (FIB 1)
    1: FIB returned 1
  0: FIB returned 2
2

where fib is just a naive recursive implementation of the Fibonacci numbers.

Using Cassette.jl, I have hacked together a similar functionality in Julia:

using Cassette

Cassette.@context TraceCtx

function indent(i::Int, s::String)
    join(fill("", i), " ") * s
end

function Cassette.overdub(ctx::TraceCtx, args...)
    needprint = :all in ctx.metadata.funs || first(args) in ctx.metadata.funs
    if needprint
        if first(args) isa Core.Builtin
            meth = "Primitive $(first(args))"
        else
            meth = "Method $(which(first(args), Base.typesof(args[2:end]...)))"
        end
        println(indent(ctx.metadata.indent, "$meth called with $(args[2:end])"))
    end
    if Cassette.canrecurse(ctx, args...)
        newctx = Cassette.similarcontext(ctx, metadata = (funs = ctx.metadata.funs, indent = ctx.metadata.indent + 4))
        res = Cassette.recurse(newctx, args...)
    else
        res = Cassette.fallback(ctx, args...)
    end
    if needprint
        println(indent(ctx.metadata.indent, "returned $res"))
    end
    res
end

macro trace(expr, funs...)
    expresc = esc(expr)
    funsesc = esc.(funs)
    :(Cassette.overdub(TraceCtx(metadata = (funs = [$(funsesc...)], indent = 0)), () -> $expresc))
end

which can be used as follows:

julia> fib(n) = if n < 2 n else fib(n-1) + fib(n-2) end
fib (generic function with 1 method)

julia> @trace fib(3) fib
   Method fib(n) in Main at REPL[6]:1 called with (3,)
       Method fib(n) in Main at REPL[6]:1 called with (2,)
           Method fib(n) in Main at REPL[6]:1 called with (1,)
           returned 1
           Method fib(n) in Main at REPL[6]:1 called with (0,)
           returned 0
       returned 1
       Method fib(n) in Main at REPL[6]:1 called with (1,)
       returned 1
   returned 2
2

The nice thing about trace is that it can be used interactively, shows all calls of selected functions no matter where and how deep they occur in the call stack and, in case of methods, also shows on which types exactly each call was dispatched.

Is this or a similar functionality already available in Julia, e.g., using some package? If not, did nobody ever wonder which methods – dispatched on which types exactly – get called across the call stack?

4 Likes

Perhaps you’d like Ghost, or is it now Umlaut?

Slightly different…

Really does not seem to be available already …
In any case, I have put together a small package from the above code snippet: TraceMe.jl

Hope some of you find it useful as well …

5 Likes

Hey @bertschi, thank you for your method tracing example!

I wanted to output a nice callgraph, so I modified it a little to output a TikzGraph graph instead:


using Cassette
using MetaGraphsNext
using TikzGraphs


Cassette.@context TraceCtx

function callstr(fun, args)
    fargs = join(args, ", ")
    "$fun($fargs)"
end

match(meth, fun::Function) = meth.name == nameof(fun)
match(meth, sym::Symbol) = sym == :all
match(meth, mod::Module) = meth.module == mod
function Cassette.overdub(ctx::TraceCtx, fun::Function, args...)
    meth = which(fun, Base.typesof(args...))
    needprint = any(match(meth, fun) for fun in ctx.metadata.funs)
	
	call_graph = ctx.metadata.call_graph
	cur_head_node = ctx.metadata.cur_head_node
	cur_v = nv(call_graph) + 1 

    if needprint
		call_graph[cur_v] = callstr(fun, args)
		call_graph[cur_head_node, cur_v] = nothing
    end
	
    if Cassette.canrecurse(ctx, fun, args...)
        newctx = Cassette.similarcontext(ctx, metadata = (
			funs = ctx.metadata.funs, 
			call_graph = call_graph, 
			cur_head_node = cur_v)
		)
        # Note: Work around potential Cassette bugs ...
        try
            res = Cassette.recurse(newctx, fun, args...)
        catch
            res = Cassette.fallback(ctx, fun, args...)
        end
    else
        res = Cassette.fallback(ctx, fun, args...)
    end
	
    if needprint
		call_graph[cur_v] = "$(callstr(fun, args)) = $res"
		call_graph[cur_head_node, cur_v] = nothing
    end
    res
end

macro trace(expr, funs...)
    expresc = esc(expr)
    funsesc = esc.(funs)

	call_graph = MetaGraph(
		DiGraph();
		label_type=Int,
		vertex_data_type=String
	)
	
	return quote
		local res = Cassette.overdub(TraceCtx(metadata = (funs = [$(funsesc...)], indent = 0, call_graph = $call_graph, cur_head_node = 1)), () -> $expresc)
		local node_labels = [string($call_graph[l]) for l in labels($call_graph)]
		local f = TikzGraphs.plot($call_graph.graph, node_labels)
		f
	end
end

and can be used the same way:

julia> fib(n) = if n < 2 n else fib(n-1) + fib(n-2) end
fib (generic function with 1 method)

julia> @trace fib(3) fib

image

This was my first delve into macros and Cassette.jl and your example was a nice starting point.

1 Like

That is pretty neat … would probably make sense to collect the information into a more general tree data structure, which could then be consumed by different backend, e.g., to produce my or your output.