summaryrefslogtreecommitdiff
path: root/tools/linkmem.py
blob: 452471d857e419f9821677bf735a453c869e1b98 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
import tempfile
import os
from os import path
import sys
import re
import subprocess
from contextlib import contextmanager
from collections import defaultdict
import colorsys

import cxxfilt
from elftools.elf.elffile import ELFFile
from elftools.elf.enums import ENUM_ST_SHNDX
from elftools.elf.descriptions import describe_symbol_type, describe_sh_type
import libarchive
import matplotlib.cm

@contextmanager
def chdir(newdir):
    old_cwd = os.getcwd()
    try:
        os.chdir(newdir)
        yield
    finally:
        os.chdir(old_cwd)

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,--print-map', *cmdline])
        lines = [ line.strip() for line in output.decode().splitlines() ]
        # FIXME also find isr vector table references

        defs = {}
        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(r'^(\.[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(r'^(\.[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(r'^(\.[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(r'^(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())
        for sym, (sec, size, obj) in defs.items():
            fn, _, member = re.match(r'^([^()]+)(\((.+)\))?$', 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()) }

                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(fr'\.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'])
                        name = refsym.name if refsym.name else elf.get_section(refsym['st_shndx']).name.split('.')[-1]
                        s.add(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

        return objs, aliases, defs, refs

@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)

hexcolor = lambda r, g, b, *_a: f'#{int(r*255):02x}{int(g*255):02x}{int(b*255):02x}'
def vhex(val):
    r,g,b,_a = matplotlib.cm.viridis(1.0-val)
    fc = hexcolor(r, g, b)  
    h,s,v = colorsys.rgb_to_hsv(r,g,b)
    cc = '#000000' if v > 0.8 else '#ffffff'
    return fc, cc

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--trace-sections', type=str, action='append', default=[])
    parser.add_argument('--trim-stubs', type=str, action='append', default=[])
    parser.add_argument('--highlight-subdirs', type=str, default=None)
    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 }
    objs, aliases, syms, refs = trace_source_files(args.linker_binary, args.linker_args, trace_sections)

    clusters = defaultdict(lambda: [])
    for sym, (sec, size, obj) in syms.items():
        clusters[obj].append((sym, sec, size))

    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 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 objs.items():
        if not isinstance(size, int):
            continue
        if obj.startswith(subdir_prefix):
            subdir_sizes[first_comp(obj[len(subdir_prefix):])] += size
        else:
            subdir_sizes['<others>'] += size

    print('Subdir sizes:', file=sys.stderr)
    for subdir, size in sorted(subdir_sizes.items(), key=lambda x: x[1]):
        print(f'{subdir:>20}: {size:>6,d} B', file=sys.stderr)

    def lookup_highlight(path):
        if args.highlight_subdirs:
            if obj.startswith(subdir_prefix):
                highlight_head = first_comp(path[len(subdir_prefix):])
                return subdir_colors[highlight_head], highlight_head
            else:
                return '#e0e0e0', None
        else:
            return '#ddf7f4', None

    with wrap('digraph G', print) as lvl1print:
        print('size="23.4,16.5!";')
        print('graph [fontsize=40];')
        print('node [fontsize=40];')
        #print('ratio="fill";')

        print('rankdir=LR;')
        print('ranksep=5;')
        print('nodesep=0.2;')
        print()

        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(objs[obj]/max_osize)
                highlight_subdir_part = f'<font face="carlito" color="{cc}" point-size="40">{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)} ({objs[obj]}B)</b></font>'
                          f'</td></tr></table>>;')
                lvl2print()
                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:
                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"];')