Add size tracing tool

1 files changed, 161 insertions, 0 deletions

diff --git a/controller/fw/tools/linkmem.py b/controller/fw/tools/linkmem.py
new file mode 100644
index 0000000..a04f31e
--- /dev/null
+++ b/controller/fw/tools/linkmem.py
@@ -0,0 +1,161 @@
+
+import tempfile
+import os
+from os import path
+import sys
+import re
+import subprocess
+from contextlib import contextmanager
+from collections import defaultdict
+
+import cxxfilt
+
+from elftools.elf.elffile import ELFFile
+from elftools.elf.descriptions import describe_symbol_type
+import libarchive
+
+@contextmanager
+def chdir(newdir):
+    old_cwd = os.getcwd()
+    try:
+        os.chdir(newdir)
+        yield
+    finally:
+        os.chdir(old_cwd)
+
+
+def trace_source_files(linker, cmdline, trace_sections=[]):
+    with tempfile.TemporaryDirectory() as tempdir:
+        out_path = path.join(tempdir, 'output.elf')
+        output = subprocess.check_output([linker, '-o', out_path, f'-Wl,--cref', *cmdline])
+        lines = [ line.strip() for line in output.decode().splitlines() ]
+        # FIXME also find isr vector table references
+
+        defs = {}
+        for line in lines[lines.index('Cross Reference Table')+3:]:
+            try:
+                *left, right = line.split()
+                if left:
+                    defs[' '.join(left)] = right 
+            except:
+                pass
+
+        refs = defaultdict(lambda: set())
+        syms = {}
+        for sym, obj in defs.items():
+            fn, _, member = re.match('^([^()]+)(\((.+)\))?$', obj).groups()
+            fn = path.abspath(fn)
+
+            if member:
+                subprocess.check_call(['ar', 'x', '--output', tempdir, fn, member])
+                fn = path.join(tempdir, member)
+
+            with open(fn, 'rb') as f:
+                elf = ELFFile(f)
+
+                symtab = elf.get_section_by_name('.symtab')
+                
+                symtab_demangled = { cxxfilt.demangle(nsym.name).replace(' ', ''): i
+                        for i, nsym in enumerate(symtab.iter_symbols()) }
+
+                def lookup_size(name):
+                    name_normalized = name.replace(' ', '')
+                    if name_normalized in symtab_demangled:
+                        entry = symtab.get_symbol(symtab_demangled[name_normalized])
+                        return entry['st_size']
+                    else:
+                        return None
+
+                syms[sym] = fn, lookup_size(sym)
+
+                s = set()
+                sec_map = { sec.name: i for i, sec in enumerate(elf.iter_sections()) }
+                sec_name = f'.rel.text.{sym}'
+                matches = [ i for name, i in sec_map.items() if re.match(f'\.rel\..*\.{sym}', name) ]
+                if matches:
+                    sec = elf.get_section(matches[0])
+                    for reloc in sec.iter_relocations():
+                        refsym = symtab.get_symbol(reloc['r_info_sym'])
+                        s.add(refsym.name)
+                        
+                        if refsym.name not in defs:
+                            syms[refsym.name] = fn, lookup_size(refsym.name)
+                refs[sym] = s
+
+                for tsec in trace_sections:
+                    matches = [ i for name, i in sec_map.items() if name == f'.rel{tsec}' ]
+                    s = set()
+                    if matches:
+                        sec = elf.get_section(matches[0])
+                        for reloc in sec.iter_relocations():
+                            refsym = symtab.get_symbol(reloc['r_info_sym'])
+                            s.add(refsym.name)
+                    refs[tsec.replace('.', '_')] |= s
+
+        syms_out = set()
+        with open(out_path, 'rb') as f:
+            elf = ELFFile(f)
+            symtab = elf.get_section_by_name('.symtab')
+            for sym in symtab.iter_symbols():
+                if describe_symbol_type(sym['st_info']['type']) == 'FUNC':
+                    syms_out.add(sym.name)
+            #for sym in defs:
+            #    entry = symtab.get_symbol_by_name(sym)
+            #    if entry is None:
+            #        syms[sym] = defs[sym], None
+            #    else:
+            #        syms[sym] = defs[sym], entry[0]['st_size']
+
+        return syms, refs, syms_out
+
+@contextmanager
+def wrap(leader='', print=print, left='{', right='}'):
+    print(leader, left)
+    yield lambda *args, **kwargs: print('   ', *args, **kwargs)
+    print(right)
+
+def mangle(name):
+    return re.sub('[^a-zA-Z0-9_]', '_', name)
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--trace-sections', type=str, action='append', default=[])
+    parser.add_argument('linker_binary')
+    parser.add_argument('linker_args', nargs=argparse.REMAINDER)
+    args = parser.parse_args()
+
+    trace_sections = args.trace_sections
+    trace_sections_mangled = { sec.replace('.', '_') for sec in trace_sections }
+    syms, refs, syms_out = trace_source_files(args.linker_binary, args.linker_args, trace_sections)
+
+    clusters = defaultdict(lambda: [])
+    for sym, (obj, size) in syms.items():
+        clusters[obj].append((sym, size))
+
+    obj_size = defaultdict(lambda: 0)
+    for name, (obj, size) in syms.items():
+        if size is not None:
+            obj_size[obj] += size
+
+    with wrap('digraph G', print) as lvl1print:
+        print('rankdir=LR;')
+        print()
+
+        for i, (obj, syms) in enumerate(clusters.items()):
+            with wrap(f'subgraph cluster_{i}', lvl1print) as lvl2print:
+                lvl2print(f'label = "{obj} <{obj_size[obj]}>";')
+                lvl2print()
+                for sym, size in syms:
+                    if sym in syms_out:
+                        lvl2print(f'{mangle(sym)}[label = "{sym} <{size}>"];')
+            lvl1print()
+
+        for start, ends in refs.items():
+            for end in ends:
+                if end and (start in syms_out or start in trace_sections_mangled) and end in syms_out:
+                    lvl1print(f'{mangle(start)} -> {mangle(end)};')
+
+        for sec in trace_sections:
+            lvl1print(f'{sec.replace(".", "_")} [label = "section {sec}"];')
+