trezor-firmware/tools/snippets/local_cache_attribute.py

"""
Helper for handling `local_cache_attribute` size optimizations in `core/src` code.

Has the possibility to transfer comments into code changes and vice-versa.

Possible improvements (TODO):
- do not stop with the renaming process when inner function is encountered
    def abc(self):
        x = self.x  # local_cache_attribute
        ghi(x)
        y = x

        def ggg():
            ...

        ghi(self.x)
        y = self.x
- do not rename when the new name is already a global symbol
    import multisig
    def abc(msg):
        multisig = msg.multisig  # local_cache_attribute
        multisig.ask(multisig)  # ERROR
- do not rename two caches with the same name
    slice_view = aprime.slice_view  # local_cache_attribute
    slice_view = bprime.slice_view  # local_cache_attribute
"""

from __future__ import annotations

import sys
from pathlib import Path

import click

try:
    import libcst as cst
    import libcst.matchers as m
except ImportError:
    click.echo("please install libcst via: pip install libcst")
    sys.exit(1)


TRANSLATED_COMMENT_MATCHER = m.SimpleStatementLine(
    body=[m.Assign()],
    trailing_whitespace=m.TrailingWhitespace(
        comment=m.Comment(value="# local_cache_attribute")
    ),
)


def attr_to_list(attr: cst.Attribute) -> list[str]:
    if m.matches(attr, m.Attribute(value=m.Name(), attr=m.Name())):
        return [attr.value.value, attr.attr.value]
    if m.matches(attr, m.Attribute(value=m.Attribute(), attr=m.Name())):
        return attr_to_list(attr.value) + [attr.attr.value]
    raise ValueError("unexpected attr format")


class Unrenamer(cst.CSTTransformer):
    def __init__(self, module: cst.Module, simplify: bool) -> None:
        self.renamers: list[tuple[cst.Name, cst.Attribute]] = []
        self.module = module
        self.simplify = simplify

    def leave_SimpleStatementLine(
        self, node: cst.SimpleStatementLine, updated: cst.CSTNode
    ) -> cst.CSTNode | None:
        if not m.matches(updated, TRANSLATED_COMMENT_MATCHER):
            return updated
        assign: cst.Assign = updated.body[0]
        name: cst.Name = assign.targets[0].target
        value_attr: cst.Attribute = assign.value
        if not isinstance(value_attr, cst.Attribute):
            raise Exception(
                f"Unexpected non-attribute assignment: {self.module.code_for_node(assign)}"
            )
        self.renamers.append((name, value_attr))

        attr_list = attr_to_list(value_attr)
        attr_str = ".".join(attr_list)
        attr_longname = "_".join(attr_list)
        orig_name = name.value

        if self.simplify and orig_name == attr_longname:
            orig_name = attr_list[-1]

        if orig_name != attr_list[-1]:
            comment_str = f"{attr_str} -> {orig_name}"
        else:
            comment_str = attr_str

        return cst.EmptyLine(
            indent=True,
            comment=cst.Comment(f"# local_cache_attribute: {comment_str}"),
        )

    def leave_Name(self, node: cst.Name, updated: cst.Name) -> cst.CSTNode:
        for old_name, attr in self.renamers:
            if updated.deep_equals(old_name):
                return attr
        return updated

    def leave_FunctionDef_body(self, node: cst.FunctionDef) -> None:
        self.renamers.clear()


class Renamer(cst.CSTTransformer):
    def __init__(self, _module: cst.Module, _simplify: bool) -> None:
        self.renamers: list[tuple[list[str], cst.Name]] = []
        self.name_is_keyword = None

    def leave_EmptyLine(
        self, node: cst.EmptyLine, updated: cst.EmptyLine
    ) -> cst.CSTNode:
        if not m.matches(node, m.EmptyLine(comment=m.Comment())):
            return updated

        comment = node.comment.value
        if not comment.startswith("# local_cache_attribute: "):
            return updated

        value_str = comment[len("# local_cache_attribute: ") :]
        if " -> " in value_str:
            value_str, target_str = value_str.split(" -> ", maxsplit=1)
        else:
            target_str = None
        attr = value_str.split(".")
        name = cst.Name(target_str or attr[-1])

        statement = cst.SimpleStatementLine(
            body=[
                cst.Assign(
                    targets=[cst.AssignTarget(target=name)],
                    value=self.process_attribute(attr),
                )
            ],
            trailing_whitespace=cst.TrailingWhitespace(
                whitespace=cst.SimpleWhitespace(value="  "),
                comment=cst.Comment(value="# local_cache_attribute"),
            ),
        )
        self.renamers.append((attr, name))
        return statement

    def visit_Name(self, node: cst.Name) -> None:
        if node is self.name_is_keyword:
            return
        for _, name in self.renamers:
            if node.deep_equals(name):
                raise Exception(f"Name {name.value} already exists in the function")

    def visit_Arg_keyword(self, node: cst.Arg) -> None:
        self.name_is_keyword = node.keyword

    def leave_Arg_keyword(self, node: cst.Arg) -> None:
        self.name_is_keyword = None

    def process_attribute(self, node: list[str]) -> cst.BaseExpression:
        assert node
        if len(node) == 1:
            return cst.Name(value=node[0])
        for old_attr, name in self.renamers:
            if node == old_attr:
                return name
        return cst.Attribute(
            value=self.process_attribute(node[:-1]), attr=cst.Name(value=node[-1])
        )

    def visit_Attribute(self, node: cst.Attribute) -> bool:
        # prevent recursing into attribute chains so that we can recurse manually
        # in leave_attribute
        return False

    def leave_Attribute(
        self, node: cst.Attribute, updated: cst.Attribute
    ) -> cst.CSTNode:
        assert node.deep_equals(updated)
        try:
            return self.process_attribute(attr_to_list(updated))
        except ValueError:
            return updated

    def leave_FunctionDef_body(self, node: cst.FunctionDef) -> None:
        self.renamers.clear()


def transform_file(
    path: Path, transformer: type[cst.CSTTransformer], simplify: bool
) -> None:
    try:
        module = cst.parse_module(path.read_text())
        modified = module.visit(transformer(module, simplify))
        if modified.code != module.code:
            path.write_text(modified.code)
            click.echo(f"Successfully converted {path}")
    except Exception as e:
        click.echo(f"Failed to convert {path}: {e}")


@click.command()
@click.argument(
    "filename", nargs=-1, type=click.Path(exists=True, file_okay=True, dir_okay=True)
)
@click.option("-r", "--reverse", is_flag=True)
@click.option("-s", "--simplify", is_flag=True)
def main(filename: list[str], reverse: bool, simplify: bool) -> None:
    if not filename:
        raise click.ClickException("No files specified")

    if reverse:
        transformer = Unrenamer
    else:
        transformer = Renamer

    for name in filename:
        path = Path(name)
        if path.is_dir():
            for subpath in path.glob("**/*.py"):
                transform_file(subpath, transformer, simplify)
        else:
            transform_file(path, transformer, simplify)


if __name__ == "__main__":
    main()