1
0
mirror of https://github.com/trezor/trezor-firmware.git synced 2024-12-29 09:38:08 +00:00
trezor-firmware/core/tools/alloc.py
2023-04-14 11:22:36 +02:00

278 lines
8.8 KiB
Python
Executable File

#!/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
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)
calls = int(calls)
filename, lineno = ident.split(":")
lineno = int(lineno)
filedata = parsed_data.setdefault(filename, {})
filedata[lineno] = {
"total_allocs": allocs,
"total_calls": calls,
"avg_allocs": allocs / calls,
}
return parsed_data
@click.group()
@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: 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: str) -> str:
if filename.startswith("src/"):
return filename[4:]
return filename
@cli.command()
@click.pass_obj
@click.argument("filename")
def annotate(obj: SharedObject, filename: str):
filename = _normalize_filename(filename)
if obj.type == "total":
def alloc_str(line: LineAllocData) -> str:
return str(line["total_allocs"])
else:
def alloc_str(line: LineAllocData) -> str:
return f"{line['avg_allocs']:.2f}"
filedata = obj.data[filename]
linedata = {lineno: alloc_str(line) for lineno, line in filedata.items()}
maxlen = max(len(l) for l in linedata.values())
lineno = 0
for line in open("src/" + filename):
lineno += 1
linecount = linedata.get(lineno, "")
print(f"{linecount:>{maxlen}} {line}", end="")
def _list(
obj: SharedObject, sort_by: str = "avg_allocs", reverse: bool = False
) -> list[tuple[str, float, int]]:
return sorted(
(
(
filename,
sum(line["avg_allocs"] for line in lines.values()),
sum(line["total_allocs"] for line in lines.values()),
)
for filename, lines in obj.data.items()
),
key=lambda x: x[1 if sort_by == "avg_allocs" else 2],
reverse=reverse,
)
@cli.command(name="list")
@click.pass_obj
@click.option("-r", "--reverse", is_flag=True)
def list_function(obj: SharedObject, reverse: bool):
if obj.type == "total":
field = "total_allocs"
def format_num(l: tuple[str, float, int]) -> str:
return f"{l[2]}"
else:
field = "avg_allocs"
def format_num(l: tuple[str, float, int]) -> str:
return f"{l[1]:.2f}"
file_sums = _list(obj, field, reverse)
maxlen = max(len(format_num(l)) for l in file_sums)
for l in file_sums:
num_str = format_num(l)
filename = l[0]
print(f"{num_str:>{maxlen}} {filename}")
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"]
return sorted(all_allocs.items(), key=lambda x: x[1], reverse=True)[:biggest_n]
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: SharedObject, htmldir: str):
file_sums = _list(obj, "total_allocs", reverse=True)
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(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:
f.write(doc.render())
if __name__ == "__main__":
cli()