1 files changed, 27 insertions, 3 deletions
diff --git a/controller/fw/tools/linkmem.py b/controller/fw/tools/linkmem.py
index 494217a..3d08d19 100644
--- a/controller/fw/tools/linkmem.py
+++ b/controller/fw/tools/linkmem.py
@@ -118,7 +118,7 @@ def wrap(leader='', print=print, left='{', right='}'):
 def mangle(name):
     return re.sub('[^a-zA-Z0-9_]', '_', name)
 
-hexcolor = lambda r, g, b: f'#{int(r*255):02x}{int(g*255):02x}{int(b*255):02x}'
+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)  
@@ -131,6 +131,7 @@ if __name__ == '__main__':
     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()
@@ -155,19 +156,42 @@ if __name__ == '__main__':
 
     max_osize = max(obj_size.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 = { path: hexcolor(*matplotlib.cm.Paired(i/len(subdir_colors))) for i, path in enumerate(subdir_colors) }
+
+    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('rankdir=LR;')
         print()
 
         for i, (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)
-                lvl2print(f'label = <<table border="0"><tr><td border="0" cellpadding="5" bgcolor="{fc}"><font face="carlito" color="{cc}"><b>{path.basename(obj)} ({obj_size[obj]}B)</b></font></td></tr></table>>;')
+                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'</td></tr></table>>;')
                 lvl2print()
                 for sym, size in syms:
                     if sym in syms_out:
                         fc, cc = vhex(size/max_size)
-                        lvl2print(f'{mangle(sym)}[label = "{sym} ({size}B)", style="rounded,filled", shape="box", fillcolor="{fc}", fontname="carlito", fontcolor="{cc}" color=None];')
+                        lvl2print(f'{mangle(sym)}[label = "{sym} ({size}B)", style="rounded,filled", shape="box", fillcolor="{fc}", fontname="carlito", fontcolor="{cc}" color=none];')
             lvl1print()
 
         for start, ends in refs.items():