nix-flame-graph.py

import sys
import functools

callstack = []

@functools.cache
def lookup(func):
    try:
        parts = func.split(":")
        with open(parts[0]) as f:
            line = f.read().split("\n")[int(parts[1])-1]
        return f"{func}: {line.strip().replace(';',';')}"
    except:
        return f"{func}: lookup failed"

def process_line(line):
    line = line.strip()
    if "function-trace" not in line:
        return
    if not line.startswith("function-trace"):
        line = line[line.find("function-trace"):]
    parts = line.split()
    mode = parts[1]
    func = lookup(parts[2])
    ts = parts[-1]
    # print(f"{mode=}, {func=}, {ts=}")
    if mode == "entered":
        callstack.append([func,ts])
    if mode == "exited":
        if callstack[-1][0] != func:
            print("invlaid callstack")
            print(f"{func=}, {callstack=}")
            exit(-1)
        start = callstack.pop()
        dt = int(ts) - int(start[1])
        if(len(callstack)):
            callstack[-1][1] = ts
        callstack_str = ";".join([frame[0] for frame in callstack])
        print(f"{callstack_str} {dt}")

if __name__ == "__main__":
    with open(sys.argv[1], 'r') as f:
        for line in f:
            process_line(line)