diff --git a/core/tools/alloc.py b/core/tools/alloc.py index 6a7aafc772..32ccc79729 100755 --- a/core/tools/alloc.py +++ b/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("") - 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("
") - def tr(self, *tds): - self.f.write("") - for td in tds: - if isinstance(td, tuple): - self.f.write(f"{td[1]}") - else: - self.f.write(f"{td}") - self.f.write("") +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("") - f.write( - f"

Total allocations: {sum(total_sum for _, _, total_sum in file_sums)}

" - ) - 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"{filename}", - ) - f.write("") + 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(" ", " "))) with open(str(htmlfile) + ".html", "w") as f: - filedata = obj.data[filename] - f.write(f"{filename}") - 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(" ", " ") - 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("") + f.write(doc.render()) if __name__ == "__main__":