TraceCalls.jl with Cassette

TraceCalls.jl

module TraceCalls
using Cassette

mutable struct Trace
    level::Int
    cutoff::Int
end

Cassette.@context TraceCtx

function Cassette.prehook(ctx::TraceCtx, args...)
    ctx.metadata.level += 1
end
function Cassette.posthook(ctx::TraceCtx, out, f, args...)
    ctx.metadata.level -= 1
    if ctx.metadata.level < ctx.metadata.cutoff
        print("  "^ctx.metadata.level)
        print(something(f, ""))
        print(something(args, "()"))
        print(" -> ")
        println(something(out, ""))
    end
    return nothing
end

trace(f, level) = Cassette.overdub(TraceCtx(metadata=Trace(0, level)), f)

macro trace(ex)
  :(trace(() -> $(esc(ex)), 2))
end

macro trace(level, ex)
  :(trace(() -> $(esc(ex)), $(esc(level))))
end

end

f(x) = x + sin(x) + cos(x)
TraceCalls.@trace 2 f(3)