function_tracer.jl

# use with Julia 1.11

using Cassette

Cassette.@context TraceCtx

mutable struct CallTrace
    depth::Int
    pending_newline::Bool
    CallTrace() = new(0, false)
end

function Cassette.prehook(ctx::TraceCtx, f, args...)
    trace = ctx.metadata

    # If there's a pending newline, print it now
    if trace.pending_newline
        println()
        trace.pending_newline = false
    end

    indent = "|  " ^ trace.depth
    print(indent, f, "(", join(args, ", "), ")")

    # Mark that we have a pending newline (might be needed if this function has nested calls)
    trace.pending_newline = true
    trace.depth += 1
    return nothing
end

function Cassette.posthook(ctx::TraceCtx, result, f, args...)
    trace = ctx.metadata
    trace.depth -= 1

    if trace.pending_newline
        # This was a leaf function - print result on same line
        println(" → ", result)
        trace.pending_newline = false
    else
        # This function had nested calls - print result on new line
        indent = "|  " ^ trace.depth
        println(indent, "→ ", result)
    end

    return nothing
end

function trace_function_calls(f, args...)
    println("=== Function Call Trace ===")
    trace = CallTrace()
    result = Cassette.overdub(TraceCtx(metadata=trace), f, args...)
    println("=== End Trace ===")
    return result
end

function mysum(x)
    s = 0
    for v in x
        s += v
    end
    return s
end


result = trace_function_calls(mysum, [1])