1#!/usr/bin/env python3
2#
3# Copyright (c) 2023 KNS Group LLC (YADRO)
4# Copyright (c) 2020 Yonatan Goldschmidt <yon.goldschmidt@gmail.com>
5#
6# SPDX-License-Identifier: Apache-2.0
7
8"""
9Stack compressor for FlameGraph
10
11This translate stack samples captured by perf subsystem into format
12used by flamegraph.pl. Translation uses .elf file to get function names
13from addresses
14
15Usage:
16    ./script/perf/stackcollapse.py <file with perf printbuf output> <ELF file>
17"""
18
19import re
20import sys
21import struct
22import binascii
23from functools import lru_cache
24from elftools.elf.elffile import ELFFile
25
26
27@lru_cache(maxsize=None)
28def addr_to_sym(addr, elf):
29    symtab = elf.get_section_by_name(".symtab")
30    for sym in symtab.iter_symbols():
31        if sym.entry.st_info.type == "STT_FUNC" and sym.entry.st_value <= addr < sym.entry.st_value + sym.entry.st_size:
32            return sym.name
33    if addr == 0:
34        return "nullptr"
35    return "[unknown]"
36
37
38def collapse(buf, elf):
39    while buf:
40        count, = struct.unpack_from(">Q", buf)
41        assert count > 0
42        addrs = struct.unpack_from(f">{count}Q", buf, 8)
43
44        func_trace = reversed(list(map(lambda a: addr_to_sym(a, elf), addrs)))
45        prev_func = next(func_trace)
46        line = prev_func
47        # merge dublicate functions
48        for func in func_trace:
49            if prev_func != func:
50                prev_func = func
51                line += ";" + func
52
53        print(line, 1)
54        buf = buf[8 + 8 * count:]
55
56
57if __name__ == "__main__":
58    elf = ELFFile(open(sys.argv[2], "rb"))
59    with open(sys.argv[1], "r") as f:
60        inp = f.read()
61
62    lines = inp.splitlines()
63    assert int(re.match(r"Perf buf length (\d+)", lines[0]).group(1)) == len(lines) - 1
64    buf = binascii.unhexlify("".join(lines[1:]))
65    collapse(buf, elf)
66