Minor fixes for nm mapping (project-chip#37873)

* Extract shorten name and add unit tests

* Fix ids

* Restyled by autopep8

* Add glob support for symbols, to quickly figure out sizes

* Update title

---------

Co-authored-by: Andrei Litvin <andreilitvin@google.com>
Co-authored-by: Restyled.io <commits@restyled.io>

Author: 3 people

scripts/tools/file_size_from_nm.py

@@ -41,6 +41,7 @@
 #    pandas
 #    plotly
 
+import fnmatch
 import logging
 import re
 import subprocess
@@ -362,6 +363,41 @@ def test_tree_display_name():
     ) == ["(anonymous namespace)", "AccessControlAttribute", "Read(args)"]
 
 
+def shorten_name(full_name: str) -> str:
+    """
+    Remove namespaces, but keep template parts
+
+    This tries to convert:
+      foo::bar::baz(int, double) -> baz(int, double)
+      foo::bar::baz<x::y>(int, double) -> baz<x::y>(int, double)
+      foo::bar::baz(some::ns::bit, double) -> baz(some::ns::bit, double)
+      foo::bar::baz<x::y>(some::ns::bit, double) -> baz<x::y>(some::ns::bit, double)
+
+    Remove all before '::', however '::' found before the first of < or (
+    """
+    limit1 = full_name.find('<')
+    limit2 = full_name.find('(')
+    if limit1 >= 0 and limit1 < limit2:
+        limit = limit1
+    else:
+        limit = limit2
+    separate_idx = full_name.rfind('::', 0, limit)
+    if separate_idx > 0:
+        short_name = full_name[separate_idx+2:]
+    else:
+        short_name = full_name
+    return short_name
+
+
+def test_shorten_name():
+    assert shorten_name("foo::bar::baz(int, double)") == "baz(int, double)"
+    assert shorten_name("foo::bar::baz<x::y>(int, double)") == "baz<x::y>(int, double)"
+    assert shorten_name("foo::bar::baz(some::ns::bit, double)") == "baz(some::ns::bit, double)"
+    assert shorten_name("foo::bar::baz<x::y>(some::ns::bit, double)") == "baz<x::y>(some::ns::bit, double)"
+    assert shorten_name("chip::app:EnabledEndpointsWithServerCluster::EnsureMatchingEndpoint()") == "EnsureMatchingEndpoint()"
+    assert shorten_name("void chip::app:EnabledEndpointsWithServerCluster::operator++()") == "operator++()"
+
+
 def build_treemap(
     name: str,
     symbols: list[Symbol],
@@ -381,7 +417,7 @@ def build_treemap(
     root = f"FILE: {name}"
     if zoom:
         root = root + f" (FILTER: {zoom})"
-    data: dict[str, list] = dict(name=[root], parent=[""], size=[0], hover=[""], name_with_size=[""], short_name=[""])
+    data: dict[str, list] = dict(name=[root], parent=[""], size=[0], hover=[""], name_with_size=[""], short_name=[""], id=[root])
 
     known_parents: set[str] = set()
     total_sizes: dict = {}
@@ -410,15 +446,19 @@ def build_treemap(
                 continue
 
         partial = ""
+        path = ""
         for name in tree_name[:-1]:
             if not partial:
                 next_value = name
             else:
                 next_value = partial + separator + name
+            parent_path = path if partial else root
+            path = path + separator + name
             if next_value not in known_parents:
                 known_parents.add(next_value)
                 data["name"].append(next_value)
-                data["parent"].append(partial if partial else root)
+                data["id"].append(path)
+                data["parent"].append(parent_path)
                 data["size"].append(0)
                 data["hover"].append(next_value)
                 data["name_with_size"].append("")
@@ -428,7 +468,8 @@ def build_treemap(
 
         # the name MUST be added
         data["name"].append(cxxfilt.demangle(symbol.name))
-        data["parent"].append(partial if partial else root)
+        data["id"].append(symbol.name)
+        data["parent"].append(path if partial else root)
         data["size"].append(symbol.size)
         data["hover"].append(f"{symbol.name} of type {symbol.symbol_type}")
         data["name_with_size"].append("")
@@ -446,27 +487,7 @@ def build_treemap(
         else:
             # When using object files, the paths hare are the full "foo::bar::....::method"
             # so clean them up a bit
-            short_name = data["short_name"][idx]
-
-            # remove namespaces, but keep template parts
-            # This tries to convert:
-            #   foo::bar::baz(int, double) -> baz(int, double)
-            #   foo::bar::baz<x::y>(int, double) -> baz<x::y>(int, double)
-            #   foo::bar::baz(some::ns:bit, double) -> baz(some::ns::bit, double)
-            #   foo::bar::baz<x::y>(some::ns:bit, double) -> baz<x::y>(some::ns::bit, double)
-            #
-            # Remove all before '::', however '::' found before the first of < or (
-            #
-            limit1 = short_name.find('<')
-            limit2 = short_name.find('(')
-            if limit1 >= 0 and limit1 < limit2:
-                limit = limit1
-            else:
-                limit = limit2
-            separate_idx = short_name.rfind('::', 0, limit)
-            if separate_idx:
-                short_name = short_name[separate_idx+2:]
-
+            short_name = shorten_name(data["short_name"][idx])
             data["name_with_size"][idx] = f"{short_name}: {data["size"][idx]}"
 
     extra_args = {}
@@ -477,7 +498,7 @@ def build_treemap(
     fig = figure_generator(
         data,
         names="name_with_size",
-        ids="name",
+        ids="id",
         parents="parent",
         values="size",
         maxdepth=max_depth,
@@ -698,14 +719,18 @@ def symbols_from_nm(elf_file: str) -> list[Symbol]:
     return symbols
 
 
-def fetch_symbols(elf_file: str, fetch: FetchStyle) -> Tuple[list[Symbol], str]:
+def fetch_symbols(elf_file: str, fetch: FetchStyle, glob_filter: Optional[str]) -> Tuple[list[Symbol], str]:
     """Returns the sumbol list and the separator used to split symbols
     """
     match fetch:
         case FetchStyle.NM:
-            return symbols_from_nm(elf_file), "::"
+            symbols, separator = symbols_from_nm(elf_file), "::"
         case FetchStyle.OBJDUMP:
-            return symbols_from_objdump(elf_file), '/'
+            symbols, separator = symbols_from_objdump(elf_file), '/'
+    if glob_filter is not None:
+        symbols = [s for s in symbols if fnmatch.fnmatch(s.name, glob_filter)]
+
+    return symbols, separator
 
 
 def list_id(tree_path: list[str]) -> str:
@@ -813,6 +838,11 @@ def compute_symbol_diff(orig: list[Symbol], base: list[Symbol]) -> list[Symbol]:
     default=None,
     help="Zoom in the graph to ONLY the specified path as root (e.g. ::chip::app)",
 )
+@click.option(
+    "--glob-filter",
+    default=None,
+    help="Glob filter by name",
+)
 @click.option(
     "--strip",
     default=None,
@@ -835,17 +865,21 @@ def main(
     zoom: Optional[str],
     strip: Optional[str],
     diff: Optional[str],
+    glob_filter: Optional[str],
 ):
     log_fmt = "%(asctime)s %(levelname)-7s %(message)s"
     coloredlogs.install(level=__LOG_LEVELS__[log_level], fmt=log_fmt)
 
-    symbols, separator = fetch_symbols(elf_file, __FETCH_STYLES__[fetch_via])
+    symbols, separator = fetch_symbols(elf_file, __FETCH_STYLES__[fetch_via], glob_filter)
     title = elf_file
 
+    if glob_filter:
+        title += f" FILTER {glob_filter}"
+
     if diff:
-        diff_symbols, _ = fetch_symbols(diff, __FETCH_STYLES__[fetch_via])
+        diff_symbols, _ = fetch_symbols(diff, __FETCH_STYLES__[fetch_via], glob_filter)
         symbols = compute_symbol_diff(symbols, diff_symbols)
-        title = f"{elf_file} COMPARED TO {diff}"
+        title += f" COMPARED TO {diff}"
 
     build_treemap(
         title, symbols, separator, __CHART_STYLES__[display_type], __COLOR_SCALES__[color], max_depth, zoom, strip