feat(core/tools): improve alloc.py for seeing allocations

[no changelog]

core/tools/alloc.py

@@ -1,14 +1,40 @@
 #!/usr/bin/env python3
+from __future__ import annotations
 
 from pathlib import Path
 from types import SimpleNamespace
+from typing import TYPE_CHECKING, Dict, Protocol, TextIO
+
+# for python37 support, is not present in typing there
+from typing_extensions import TypedDict
+
 import click
+from dominate import document
+from dominate.tags import *
+from dominate.util import raw
 
 HERE = Path(__file__).resolve().parent
 
 
-def parse_alloc_data(alloc_data):
-    parsed_data = {}
+if TYPE_CHECKING:
+
+    class LineAllocData(TypedDict):
+        total_allocs: int
+        total_calls: int
+        avg_allocs: float
+
+    # {filename:{lineno:LineAllocData}}
+    alloc_data_dict = Dict[str, Dict[int, LineAllocData]]
+
+    class SharedObject(Protocol):
+        data: alloc_data_dict
+        type: str
+
+
+def parse_alloc_data(
+    alloc_data: TextIO,
+) -> alloc_data_dict:
+    parsed_data: alloc_data_dict = {}
     for line in alloc_data:
         ident, allocs, calls = line.strip().split(" ")
         allocs = int(allocs)
@@ -29,11 +55,14 @@ def parse_alloc_data(alloc_data):
 @click.pass_context
 @click.option("-a", "--alloc-data", type=click.File(), default="src/alloc_data.txt")
 @click.option("-t", "--type", type=click.Choice(("total", "avg")), default="avg")
-def cli(ctx, alloc_data, type):
-    ctx.obj = SimpleNamespace(data=parse_alloc_data(alloc_data), type=type)
+def cli(ctx: click.Context, alloc_data: TextIO, type: str):
+    shared_obj: SharedObject = SimpleNamespace()  # type: ignore
+    shared_obj.data = parse_alloc_data(alloc_data)
+    shared_obj.type = type
+    ctx.obj = shared_obj
 
 
-def _normalize_filename(filename):
+def _normalize_filename(filename: str) -> str:
     if filename.startswith("src/"):
         return filename[4:]
     return filename
@@ -42,13 +71,18 @@ def _normalize_filename(filename):
 @cli.command()
 @click.pass_obj
 @click.argument("filename")
-def annotate(obj, filename):
+def annotate(obj: SharedObject, filename: str):
     filename = _normalize_filename(filename)
 
     if obj.type == "total":
-        alloc_str = lambda line: str(line["total_allocs"])
+
+        def alloc_str(line: LineAllocData) -> str:
+            return str(line["total_allocs"])
+
     else:
-        alloc_str = lambda line: f"{line['avg_allocs']:.2f}"
+
+        def alloc_str(line: LineAllocData) -> str:
+            return f"{line['avg_allocs']:.2f}"
 
     filedata = obj.data[filename]
 
@@ -62,7 +96,9 @@ def annotate(obj, filename):
         print(f"{linecount:>{maxlen}}  {line}", end="")
 
 
-def _list(obj, sort_by="avg_allocs", reverse=False):
+def _list(
+    obj: SharedObject, sort_by: str = "avg_allocs", reverse: bool = False
+) -> list[tuple[str, float, int]]:
     return sorted(
         (
             (
@@ -77,16 +113,21 @@ def _list(obj, sort_by="avg_allocs", reverse=False):
     )
 
 
-@cli.command()
+@cli.command(name="list")
 @click.pass_obj
 @click.option("-r", "--reverse", is_flag=True)
-def list(obj, reverse):
+def list_function(obj: SharedObject, reverse: bool):
     if obj.type == "total":
         field = "total_allocs"
-        format_num = lambda l: f"{l[2]}"
+
+        def format_num(l: tuple[str, float, int]) -> str:
+            return f"{l[2]}"
+
     else:
         field = "avg_allocs"
-        format_num = lambda l: f"{l[1]:.2f}"
+
+        def format_num(l: tuple[str, float, int]) -> str:
+            return f"{l[1]:.2f}"
 
     file_sums = _list(obj, field, reverse)
 
@@ -97,77 +138,139 @@ def list(obj, reverse):
         print(f"{num_str:>{maxlen}}  {filename}")
 
 
-class HtmlTable:
-    def __init__(self, f):
-        self.f = f
+def get_biggest_line_allocations(
+    obj: SharedObject, biggest_n: int
+) -> list[tuple[str, float]]:
+    all_allocs: dict[str, float] = {}
+    for file, line_stats in obj.data.items():
+        for line, stats in line_stats.items():
+            all_allocs[f"{file}:{line}"] = stats["avg_allocs"]
 
-    def __enter__(self):
-        self.f.write("<table>")
-        return self
+    return sorted(all_allocs.items(), key=lambda x: x[1], reverse=True)[:biggest_n]
 
-    def __exit__(self, type, value, traceback):
-        self.f.write("</table>")
 
-    def tr(self, *tds):
-        self.f.write("<tr>")
-        for td in tds:
-            if isinstance(td, tuple):
-                self.f.write(f"<td {td[0]}><tt>{td[1]}</tt></td>")
-            else:
-                self.f.write(f"<td><tt>{td}</tt></td>")
-        self.f.write("</tr>")
+def get_biggest_n_lines_for_each_file(
+    obj: SharedObject, biggest_n: int
+) -> dict[str, list[int]]:
+    biggest_file_allocs: dict[str, list[int]] = {}
+    for file, line_stats in obj.data.items():
+        biggest = sorted(
+            line_stats.items(), key=lambda x: x[1]["avg_allocs"], reverse=True
+        )[:biggest_n]
+        biggest_file_allocs[file] = [line for line, _stats in biggest]
+    return biggest_file_allocs
 
 
 @cli.command()
 @click.pass_obj
 @click.argument("htmldir")
-def html(obj, htmldir):
+def html(obj: SharedObject, htmldir: str):
     file_sums = _list(obj, "total_allocs", reverse=True)
-    style_grey = "style='color: grey'"
-    style_right = "style='text-align: right'"
+    style_grey = "color: grey"
+    style_red = "color: red;"
+    style_blue = "color: blue;"
+    style_right = "text-align: right;"
+    css_smaller_mono = (
+        "body { font-size: 80%; font-family: 'Courier New', Courier, monospace; }"
+    )
 
+    n_biggest = 50
+    biggest_lines = get_biggest_line_allocations(obj, n_biggest)
+    for location, avg_alloc in reversed(biggest_lines):
+        # Prepending core/src so it can be opened via alt+click in VSCode
+        print(f"{avg_alloc:.2f} core/src/{location}")
+
+    # Create index.html - two tables
+    doc = document(title="Firmware allocations")
+    with doc.head:
+        meta(charset="utf-8")
+        style(css_smaller_mono)
+    with doc:
+        h3(f"{n_biggest} biggest allocations")
+        with table():
+            with thead():
+                with tr():
+                    th("alloc", style=style_right)
+                    th("file:line")
+            with tbody():
+                for location, avg_alloc in biggest_lines:
+                    filename, lineno = location.split(":")
+                    with tr():
+                        td(f"{avg_alloc:.2f}", style=style_right)
+                        td(
+                            a(
+                                location,
+                                href=f"{filename}.html#{lineno}",
+                                target="_blank",
+                            )
+                        )
+        h3(f"Total allocations: {sum(total_sum for _, _, total_sum in file_sums)}")
+        with table():
+            with thead():
+                with tr():
+                    th("avg", style=style_right)
+                    th("total", style=style_right)
+                    th("file")
+            with tbody():
+                for filename, avg_sum, total_sum in file_sums:
+                    with tr():
+                        td(f"{avg_sum:.2f}", style=style_right)
+                        td(total_sum, style=style_right)
+                        td(
+                            a(
+                                filename,
+                                href=f"{filename}.html",
+                                target="_blank",
+                            )
+                        )
     with open(f"{htmldir}/index.html", "w") as f:
-        f.write("<html>")
-        f.write(
-            f"<h3>Total allocations: {sum(total_sum for _, _, total_sum in file_sums)}</h3>"
-        )
-        with HtmlTable(f) as table:
-            table.tr((style_right, "avg"), (style_right, "total"), "")
-            for filename, avg_sum, total_sum in file_sums:
-                table.tr(
-                    (style_right, f"{avg_sum:.2f}"),
-                    (style_right, total_sum),
-                    f"<a href='{filename}.html'>{filename}</a>",
-                )
-        f.write("</html>")
+        f.write(doc.render())
 
+    # So we can highlight biggest allocations in each file
+    biggest_n_lines_for_each_file = get_biggest_n_lines_for_each_file(obj, 5)
+
+    # Create HTML for each file - one table in each
     for filename in file_sums:
         filename = _normalize_filename(filename[0])
         htmlfile = Path(htmldir) / filename
         htmlfile.parent.mkdir(parents=True, exist_ok=True)
 
+        doc = document(title=filename)
+        with doc.head:
+            meta(charset="utf-8")
+            style(css_smaller_mono)
+        with doc:
+            with table():
+                with thead():
+                    with tr():
+                        th("#", style=style_grey)
+                        th("avg", style=style_right)
+                        th("total", style=style_right)
+                        th("")
+                with tbody():
+                    lineno = 0
+                    for line in open(HERE.parent / "src" / filename):
+                        lineno += 1
+                        line_info = obj.data[filename].get(lineno, {})
+                        total = line_info.get("total_allocs", 0)
+                        avg = line_info.get("avg_allocs", 0)
+
+                        if lineno in biggest_n_lines_for_each_file[filename]:
+                            row_style = style_red
+                        elif avg > 0:
+                            row_style = style_blue
+                        else:
+                            row_style = None
+
+                        with tr(style=row_style, id=lineno):
+                            td(lineno, style=style_grey)
+                            td(f"{avg:.2f}", style=style_right)
+                            td(total, style=style_right)
+                            # Creating nonbreaking space, otherwise prefix
+                            # whitespace is stripped
+                            td(raw(line.rstrip("\n").replace(" ", "&nbsp;")))
         with open(str(htmlfile) + ".html", "w") as f:
-            filedata = obj.data[filename]
-            f.write(f"<html><title>{filename}</title>")
-            with HtmlTable(f) as table:
-                table.tr(
-                    (style_grey, "#"), (style_right, "avg"), (style_right, "total"), ""
-                )
-
-                lineno = 0
-                for line in open(HERE.parent / "src" / filename):
-                    line = line.rstrip("\n").replace(" ", "&nbsp;")
-                    lineno += 1
-                    total = filedata.get(lineno, {}).get("total_allocs", 0)
-                    avg = filedata.get(lineno, {}).get("avg_allocs", 0)
-
-                    table.tr(
-                        (style_grey, lineno),
-                        (style_right, f"{avg:.2f}"),
-                        (style_right, total),
-                        line,
-                    )
-            f.write("</html>")
+            f.write(doc.render())
 
 
 if __name__ == "__main__":