Fix up linkmem script

2 files changed, 117 insertions, 73 deletions

diff --git a/controller/fw/Makefile b/controller/fw/Makefile
index 0faea30..11b68b8 100644
--- a/controller/fw/Makefile
+++ b/controller/fw/Makefile
@@ -306,6 +306,7 @@ clean:
 	rm -rf $(BUILDDIR)/src
 	rm -rf $(BUILDDIR)/generated
 	rm -f $(BUILDDIR)/$(BINARY)
+	rm -f $(BUILDDIR)/$(BINARY:.elf=.map)
 	rm -f $(BUILDDIR)/$(BINARY:.elf=-symbol-sizes.dot)
 	rm -f $(BUILDDIR)/$(BINARY:.elf=-symbol-sizes.pdf)
 	rm -f $(BUILDDIR)/tools/freq_meas_test
diff --git a/controller/fw/tools/linkmem.py b/controller/fw/tools/linkmem.py
index db8ddb8..fbd6dde 100644
--- a/controller/fw/tools/linkmem.py
+++ b/controller/fw/tools/linkmem.py
@@ -25,26 +25,106 @@ def chdir(newdir):
     finally:
         os.chdir(old_cwd)
 
-
-def trace_source_files(linker, cmdline, trace_sections=[]):
+def keep_last(it, first=None):
+    last = first
+    for elem in it:
+        yield last, elem
+        last = elem
+
+def delim(start, end, it, first_only=True):
+    found = False
+    for elem in it:
+        if end(elem):
+            if first_only:
+                return
+            found = False
+        elif start(elem):
+            found = True
+        elif found:
+            yield elem
+
+def delim_prefix(start, end, it):
+    yield from delim(lambda l: l.startswith(start), lambda l: end is not None and l.startswith(end), it)
+
+def trace_source_files(linker, cmdline, trace_sections=[], total_sections=['.text', '.data', '.rodata']):
     with tempfile.TemporaryDirectory() as tempdir:
         out_path = path.join(tempdir, 'output.elf')
-        output = subprocess.check_output([linker, '-o', out_path, f'-Wl,--cref', *cmdline])
+        output = subprocess.check_output([linker, '-o', out_path, f'-Wl,--print-map', *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
+        objs = defaultdict(lambda: 0)
+        aliases = {}
+        sec_name = None
+        last_loc = None
+        last_sym = None
+        line_cont = None
+        for last_line, line in keep_last(delim_prefix('Linker script and memory map', 'OUTPUT', lines), first=''):
+            if not line or line.startswith('LOAD '):
+                sec_name = None
+                continue
+
+            # first part of continuation line
+            if m := re.match('^(\.[0-9a-zA-Z-_.]+)$', line):
+                line_cont = line
+                sec_name = None
+                continue
+
+            if line_cont:
+                line = line_cont + ' ' + line
+                line_cont = None
+
+            # -ffunction-sections/-fdata-sections section
+            if m := re.match('^(\.[0-9a-zA-Z-_.]+)\.([0-9a-zA-Z-_.]+)\s+(0x[0-9a-f]+)\s+(0x[0-9a-f]+)\s+(\S+)$', line):
+                sec, sym, loc, size, obj = m.groups()
+                *_, sym = sym.rpartition('.')
+                sym = cxxfilt.demangle(sym)
+                size = int(size, 16)
+                obj = path.abspath(obj)
+
+                if sec not in total_sections:
+                    size = 0
+
+                objs[obj] += size
+                defs[sym] = (sec, size, obj)
+
+                sec_name, last_loc, last_sym = sec, loc, sym
+                continue
+            
+            # regular (no -ffunction-sections/-fdata-sections) section
+            if m := re.match('^(\.[0-9a-zA-Z-_]+)\s+(0x[0-9a-f]+)\s+(0x[0-9a-f]+)\s+(\S+)$', line):
+                sec, _loc, size, obj = m.groups()
+                size = int(size, 16)
+                obj = path.abspath(obj)
+
+                if sec in total_sections:
+                    objs[obj] += size
+
+                sec_name = sec
+                last_loc, last_sym = None, None
+                continue
+            
+            # symbol def
+            if m := re.match('^(0x[0-9a-f]+)\s+(\S+)$', line):
+                loc, sym = m.groups()
+                sym = cxxfilt.demangle(sym)
+                loc = int(loc, 16)
+                if sym in defs:
+                    continue
+
+                if loc == last_loc:
+                    assert last_sym is not None
+                    aliases[sym] = last_sym
+                else:
+                    assert sec_name
+                    defs[sym] = (sec_name, None, obj)
+                    last_loc, last_sym = loc, sym
+
+                continue
 
         refs = defaultdict(lambda: set())
-        syms = {}
-        for sym, obj in defs.items():
+        for sym, (sec, size, obj) in defs.items():
             fn, _, member = re.match('^([^()]+)(\((.+)\))?$', obj).groups()
             fn = path.abspath(fn)
 
@@ -60,22 +140,6 @@ def trace_source_files(linker, cmdline, trace_sections=[]):
                 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])
-                        index = entry['st_shndx']
-                        if index in ENUM_ST_SHNDX:
-                            return 'no bits'
-                        sec = elf.get_section(index)
-                        if describe_sh_type(sec['sh_type']) == 'NOBITS':
-                            return 'no bits'
-                        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()) }
                 matches = [ i for name, i in sec_map.items() if re.match(f'\.rel\..*\.{sym}', name) ]
@@ -85,9 +149,6 @@ def trace_source_files(linker, cmdline, trace_sections=[]):
                         refsym = symtab.get_symbol(reloc['r_info_sym'])
                         name = refsym.name if refsym.name else elf.get_section(refsym['st_shndx']).name.split('.')[-1]
                         s.add(name)
-                        
-                        if name not in defs:
-                            syms[name] = fn, lookup_size(name)
                 refs[sym] = s
 
                 for tsec in trace_sections:
@@ -100,21 +161,7 @@ def trace_source_files(linker, cmdline, trace_sections=[]):
                             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']) in ('FUNC', 'OBJECT'):
-                    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
+        return objs, aliases, defs, refs
 
 @contextmanager
 def wrap(leader='', print=print, left='{', right='}'):
@@ -145,32 +192,22 @@ if __name__ == '__main__':
 
     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)
-
-    for name, (obj, size) in syms.items():
-        if path.basename(obj) in args.trim_stubs and (not isinstance(size, int) or size <= 8) and not refs.get(name):
-            syms_out.discard(name)
+    objs, aliases, syms, refs = trace_source_files(args.linker_binary, args.linker_args, trace_sections)
 
     clusters = defaultdict(lambda: [])
-    for sym, (obj, size) in syms.items():
-        if sym in syms_out:
-            clusters[obj].append((sym, size))
+    for sym, (sec, size, obj) in syms.items():
+        clusters[obj].append((sym, sec, size))
 
-    obj_size = defaultdict(lambda: 0)
-    for sym, (obj, size) in syms.items():
-        if isinstance(size, int) and sym in syms_out:
-            obj_size[obj] += size
-    max_size = max([ size for _obj, size in syms.values() if isinstance(size, int) ])
-
-    max_osize = max(obj_size.values())
+    max_ssize = max(size or 0 for _sec, size, _obj in syms.values())
+    max_osize = max(objs.values())
 
     subdir_prefix = path.abspath(args.highlight_subdirs) + '/' if args.highlight_subdirs else '### NO HIGHLIGHT ###'
     first_comp = lambda le_path: path.dirname(le_path).partition(os.sep)[0]
-    subdir_colors = sorted({ first_comp(obj[len(subdir_prefix):]) for obj in obj_size if obj.startswith(subdir_prefix) })
+    subdir_colors = sorted({ first_comp(obj[len(subdir_prefix):]) for obj in objs if obj.startswith(subdir_prefix) })
     subdir_colors = { path: hexcolor(*matplotlib.cm.Pastel1(i/len(subdir_colors))) for i, path in enumerate(subdir_colors) }
 
     subdir_sizes = defaultdict(lambda: 0)
-    for obj, size in obj_size.items():
+    for obj, size in objs.items():
         if not isinstance(size, int):
             continue
         if obj.startswith(subdir_prefix):
@@ -198,30 +235,36 @@ if __name__ == '__main__':
         print('nodesep=0.2;')
         print()
 
-        for i, (obj, syms) in enumerate(clusters.items()):
+        for i, (obj, obj_syms) in enumerate(clusters.items()):
             with wrap(f'subgraph cluster_{i}', lvl1print) as lvl2print:
                 print('style = "filled";')
                 highlight_color, highlight_head = lookup_highlight(obj)
                 print(f'bgcolor = "{highlight_color}";')
                 print('pencolor = none;')
-                fc, cc = vhex(obj_size[obj]/max_osize)
+                fc, cc = vhex(objs[obj]/max_osize)
                 highlight_subdir_part = f'<font face="carlito" color="{cc}" point-size="12">{highlight_head} / </font>' if highlight_head else ''
                 lvl2print(f'label = <<table border="0"><tr><td border="0" cellpadding="5" bgcolor="{fc}">'
                           f'{highlight_subdir_part}'
-                          f'<font face="carlito" color="{cc}"><b>{path.basename(obj)} ({obj_size[obj]}B)</b></font>'
+                          f'<font face="carlito" color="{cc}"><b>{path.basename(obj)} ({objs[obj]}B)</b></font>'
                           f'</td></tr></table>>;')
                 lvl2print()
-                for sym, size in syms:
-                    if sym in syms_out:
-                        size_s = f'{size}B' if isinstance(size, int) else size
-                        fc, cc = vhex(size/max_size if isinstance(size, int) else 0)
-                        lvl2print(f'{mangle(sym)}[label = "{sym} ({size_s})", style="rounded,filled", shape="box", fillcolor="{fc}", fontname="carlito", fontcolor="{cc}" color=none];')
+                for sym, sec, size in obj_syms:
+                    has_size = isinstance(size, int) and size > 0
+                    size_s = f' ({size}B)' if has_size else ''
+                    fc, cc = vhex(size/max_ssize) if has_size else ('#ffffff', '#000000')
+                    shape = 'box' if sec == '.text' else 'oval'
+                    lvl2print(f'{mangle(sym)}[label = "{sym}{size_s}", style="rounded,filled", shape="{shape}", fillcolor="{fc}", fontname="carlito", fontcolor="{cc}" color=none];')
             lvl1print()
 
+        edges = set()
         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)} [style="bold", color="#333333"];')
+                end = aliases.get(end, end)
+                if (start in syms or start in trace_sections_mangled) and end in syms:
+                    edges.add((start, end))
+
+        for start, end in edges:
+            lvl1print(f'{mangle(start)} -> {mangle(end)} [style="bold", color="#333333"];')
 
         for sec in trace_sections:
             lvl1print(f'{sec.replace(".", "_")} [label = "section {sec}", shape="box", style="filled,bold"];')