path: root/build-aux/measurestack/analyze.py

# build-aux/measurestack/analyze.py - Analyze stack sizes for compiled objects
#
# Copyright (C) 2024-2025  Luke T. Shumaker <lukeshu@lukeshu.com>
# SPDX-License-Identifier: AGPL-3.0-or-later

import random
import re
import sys
import typing

from . import vcg

# Whether to print "//dbg-cache:" on cache writes
dbg_cache = False
# Whether to print the graph in a /* comment */ before processing it
dbg_dumpgraph = False
# Whether to print "//dbg-nstatic:" lines that trace nstatic() execution
dbg_nstatic = False
# Whether to disable nstatic() caching (but does NOT disable any cache-related debug logging)
dbg_nocache = False
# Whether to sort things for consistently-ordered execution, or shuffle things to detect bugs
dbg_sort: typing.Literal["unsorted", "sorted", "shuffled"] = "unsorted"

# pylint: disable=unused-variable
__all__ = [
    "BaseName",
    "QName",
    "UsageKind",
    "Node",
    "maybe_sorted",
    "AnalyzeResultVal",
    "AnalyzeResultGroup",
    "AnalyzeResult",
    "analyze",
]


def dumps(x: typing.Any, depth: int = 0, compact: bool = False) -> str:
    match x:
        case int() | str() | None:
            return repr(x)
        case dict():
            if len(x) == 0:
                return "{}"
            ret = "{"
            if not compact:
                ret += "\n"
            for k, v in x.items():
                if not compact:
                    ret += "\t" * (depth + 1)
                ret += dumps(k, depth + 1, True)
                ret += ":"
                if not compact:
                    ret += " "
                ret += dumps(v, depth + 1, compact)
                ret += ","
                if not compact:
                    ret += "\n"
            if not compact:
                ret += "\t" * depth
            ret += "}"
            return ret
        case list():
            if len(x) == 0:
                return "[]"
            ret = "["
            if not compact:
                ret += "\n"
            for v in x:
                if not compact:
                    ret += "\t" * (depth + 1)
                ret += dumps(v, depth + 1, compact)
                ret += ","
                if not compact:
                    ret += "\n"
            if not compact:
                ret += "\t" * depth
            ret += "]"
            return ret
        case set():
            if len(x) == 0:
                return "set()"
            ret = "{"
            if not compact:
                ret += "\n"
            for v in x:
                if not compact:
                    ret += "\t" * (depth + 1)
                ret += dumps(v, depth + 1, compact)
                ret += ","
                if not compact:
                    ret += "\n"
            if not compact:
                ret += "\t" * depth
            ret += "}"
            return ret
        case _:
            if hasattr(x, "__dict__"):
                return f"{x.__class__.__name__}(*{dumps(x.__dict__, depth, compact)})"
            return f"TODO({x.__class__.__name__})"


# types ########################################################################


class BaseName:
    # class ##########################################################

    _interned: dict[str, "BaseName"] = {}

    def __new__(cls, content: str) -> "BaseName":
        if ":" in content:
            raise ValueError(f"invalid non-qualified name: {content!r}")
        content = sys.intern(content)
        if content not in cls._interned:
            self = super().__new__(cls)
            self._content = content
            cls._interned[content] = self
        return cls._interned[content]

    # instance #######################################################

    _content: str

    def __str__(self) -> str:
        return self._content

    def __repr__(self) -> str:
        return f"BaseName({self._content!r})"

    def __format__(self, fmt_spec: str, /) -> str:
        return repr(self)

    def __eq__(self, other: typing.Any) -> bool:
        assert isinstance(
            other, BaseName
        ), f"comparing BaseName with {other.__class__.__name__}"
        return self._content == other._content

    def __lt__(self, other: "BaseName") -> bool:
        return self._content < other._content

    def __hash__(self) -> int:
        return self._content.__hash__()

    def as_qname(self) -> "QName":
        return QName(self._content)


class QName:
    # class ##########################################################

    _interned: dict[str, "QName"] = {}

    def __new__(cls, content: str) -> "QName":
        content = sys.intern(content)
        if content not in cls._interned:
            self = super().__new__(cls)
            self._content = content
            self._base = None
            cls._interned[content] = self
        return cls._interned[content]

    # instance #######################################################

    _content: str
    _base: BaseName | None

    def __str__(self) -> str:
        return self._content

    def __repr__(self) -> str:
        return f"QName({self._content!r})"

    def __format__(self, fmt_spec: str, /) -> str:
        return repr(self)

    def __eq__(self, other: typing.Any) -> bool:
        assert isinstance(
            other, QName
        ), f"comparing QName with {other.__class__.__name__}"
        return self._content == other._content

    def __lt__(self, other: "QName") -> bool:
        return self._content < other._content

    def __hash__(self) -> int:
        return self._content.__hash__()

    def base(self) -> BaseName:
        if self._base is None:
            self._base = BaseName(self._content.rsplit(":", 1)[-1].split(".", 1)[0])
        return self._base


UsageKind: typing.TypeAlias = typing.Literal["static", "dynamic", "dynamic,bounded"]


class Node:
    # from .title (`static` and `__weak` functions are prefixed with
    # the compilation unit .c file.  For static functions that's fine,
    # but we'll have to handle it specially for __weak.).
    funcname: QName
    # .label is "{funcname}\n{location}\n{nstatic} bytes (static}\n{ndynamic} dynamic objects"
    location: str
    usage_kind: UsageKind
    nstatic: int
    ndynamic: int

    # edges with .sourcename set to this node, val is if it's
    # OK/expected that the function be missing.
    calls: dict[QName, bool]


class AnalyzeResultVal(typing.NamedTuple):
    nstatic: int
    cnt: int


class AnalyzeResultGroup(typing.NamedTuple):
    rows: dict[QName, AnalyzeResultVal]


class AnalyzeResult(typing.NamedTuple):
    groups: dict[str, AnalyzeResultGroup]
    missing: set[QName]
    dynamic: set[QName]

    included_funcs: set[QName]


class SkipModel(typing.NamedTuple):
    """Running the skipmodel calls `.fn(chain, ...)` with the chain
    consisting of the last few items of the input chain.

    If `.nchain` is an int:

     - the chain is the last `.nchain` items or the input chain.  If
       the input chain is not that long, then `.fn` is not called and
       the call is *not* skipped.

    If `.nchain` is a collection:

     - the chain starts with the *last* occurance of `.nchain` in the
       input chain.  If the input chain does not contain a member of
       the collection, then .fn is called with an empty chain.
    """

    nchain: int | typing.Collection[BaseName]
    fn: typing.Callable[[typing.Sequence[QName], Node, QName], bool]

    def __call__(
        self, chain: typing.Sequence[QName], node: Node, call: QName
    ) -> tuple[bool, int]:
        match self.nchain:
            case int():
                if len(chain) >= self.nchain:
                    _chain = chain[-self.nchain :]
                    return self.fn(_chain, node, call), len(_chain) + 1
                return False, 0
            case _:
                for i in reversed(range(len(chain))):
                    if chain[i].base() in self.nchain:
                        _chain = chain[i:]
                        return self.fn(_chain, node, call), len(_chain) + 1
                return self.fn([], node, call), 1


class Application(typing.Protocol):
    def extra_nodes(self) -> typing.Collection[Node]: ...
    def indirect_callees(
        self, elem: vcg.VCGElem
    ) -> tuple[typing.Collection[QName], bool]: ...
    def skipmodels(self) -> dict[BaseName, SkipModel]: ...


# code #########################################################################

re_node_label = re.compile(
    r"(?P<funcname>[^\n]+)\n"
    + r"(?P<location>[^\n]+:[0-9]+:[0-9]+)\n"
    + r"(?P<nstatic>[0-9]+) bytes \((?P<usage_kind>static|dynamic|dynamic,bounded)\)\n"
    + r"(?P<ndynamic>[0-9]+) dynamic objects"
    + r"(?:\n.*)*",
    flags=re.MULTILINE,
)


class _Graph:
    graph: dict[QName, Node]
    qualified: dict[BaseName, QName]

    _resolve_cache: dict[QName, QName | None]

    def __init__(self) -> None:
        self._resolve_cache = {}

    def _resolve_funcname(self, funcname: QName) -> QName | None:
        s = str(funcname)
        is_qualified = ":" in s

        # Handle `ld --wrap` functions
        if not is_qualified:
            with_wrap = QName(f"__wrap_{s}")
            if with_wrap in self.graph:
                return with_wrap
            if s.startswith("__real_"):
                without_real = QName(s[len("__real_") :])
                if without_real in self.graph:
                    funcname = without_real

        # Usual case
        if funcname in self.graph:
            return funcname

        # Handle `__weak`/`[[gnu::weak]]` functions
        if not is_qualified:
            return self.qualified.get(BaseName(s))

        return None

    def resolve_funcname(self, funcname: QName) -> QName | None:
        if funcname not in self._resolve_cache:
            self._resolve_cache[funcname] = self._resolve_funcname(funcname)
        return self._resolve_cache[funcname]


if typing.TYPE_CHECKING:
    from _typeshed import SupportsRichComparisonT as _T_sortable

_T = typing.TypeVar("_T")


@typing.overload
def maybe_sorted(
    unsorted: typing.Iterable["_T_sortable"], /, *, key: None = None
) -> typing.Iterable["_T_sortable"]: ...
@typing.overload
def maybe_sorted(
    unsorted: typing.Iterable[_T], /, *, key: typing.Callable[[_T], "_T_sortable"]
) -> typing.Iterable[_T]: ...


def maybe_sorted(
    unsorted: typing.Iterable[_T],
    /,
    *,
    key: typing.Callable[[_T], "_T_sortable"] | None = None,
) -> typing.Iterable[_T]:
    match dbg_sort:
        case "unsorted":
            return unsorted
        case "sorted":
            return sorted(unsorted, key=key)  # type: ignore
        case "shuffled":
            ret = [*unsorted]
            random.shuffle(ret)
            return ret


def _make_graph(
    ci_fnames: typing.Collection[str],
    app: Application,
) -> _Graph:
    graph: dict[QName, Node] = {}
    qualified: dict[BaseName, set[QName]] = {}

    def handle_elem(elem: vcg.VCGElem) -> None:
        match elem.typ:
            case "node":
                node = Node()
                node.calls = {}
                skip = False
                for k, v in elem.attrs.items():
                    match k:
                        case "title":
                            node.funcname = QName(v)
                        case "label":
                            if elem.attrs.get("shape", "") != "ellipse":
                                m = re_node_label.fullmatch(v)
                                if not m:
                                    raise ValueError(f"unexpected label value {v!r}")
                                node.location = m.group("location")
                                node.usage_kind = typing.cast(
                                    UsageKind, m.group("usage_kind")
                                )
                                node.nstatic = int(m.group("nstatic"))
                                node.ndynamic = int(m.group("ndynamic"))
                        case "shape":
                            if v != "ellipse":
                                raise ValueError(f"unexpected shape value {v!r}")
                            skip = True
                        case _:
                            raise ValueError(f"unknown edge key {k!r}")
                if not skip:
                    if node.funcname in graph:
                        raise ValueError(f"duplicate node {node.funcname}")
                    graph[node.funcname] = node
                    if ":" in str(node.funcname):
                        basename = node.funcname.base()
                        if basename not in qualified:
                            qualified[basename] = set()
                        qualified[basename].add(node.funcname)
            case "edge":
                caller: QName | None = None
                callee: QName | None = None
                for k, v in elem.attrs.items():
                    match k:
                        case "sourcename":
                            caller = QName(v)
                        case "targetname":
                            callee = QName(v)
                        case "label":
                            pass
                        case _:
                            raise ValueError(f"unknown edge key {k!r}")
                if caller is None or callee is None:
                    raise ValueError(f"incomplete edge: {elem.attrs!r}")
                if caller not in graph:
                    raise ValueError(f"unknown caller: {caller}")
                if callee == QName("__indirect_call"):
                    callees, missing_ok = app.indirect_callees(elem)
                    for callee in maybe_sorted(callees):
                        if callee not in graph[caller].calls:
                            graph[caller].calls[callee] = missing_ok
                else:
                    graph[caller].calls[callee] = False
            case _:
                raise ValueError(f"unknown elem type {elem.typ!r}")

    for ci_fname in maybe_sorted(ci_fnames):
        with open(ci_fname, "r", encoding="utf-8") as fh:
            for elem in vcg.parse_vcg(fh):
                handle_elem(elem)

    def sort_key(node: Node) -> QName:
        return node.funcname

    for node in maybe_sorted(app.extra_nodes(), key=sort_key):
        if node.funcname in graph:
            raise ValueError(f"duplicate node {node.funcname}")
        graph[node.funcname] = node

    ret = _Graph()
    ret.graph = graph
    ret.qualified = {}
    for bname, qnames in qualified.items():
        if len(qnames) == 1:
            ret.qualified[bname] = next(name for name in qnames)
    return ret


def analyze(
    *,
    ci_fnames: typing.Collection[str],
    app_func_filters: dict[str, typing.Callable[[QName], tuple[int, bool]]],
    app: Application,
    cfg_max_call_depth: int,
) -> AnalyzeResult:
    graphdata = _make_graph(ci_fnames, app)
    if dbg_dumpgraph:
        print(f"/* {dumps(graphdata)} */")

    missing: set[QName] = set()
    dynamic: set[QName] = set()
    included_funcs: set[QName] = set()

    track_inclusion: bool = True

    skipmodels = app.skipmodels()
    for name, model in skipmodels.items():
        if not isinstance(model.nchain, int):
            assert len(model.nchain) > 0

    _nstatic_cache: dict[QName, int] = {}

    def _nstatic(chain: list[QName], funcname: QName) -> tuple[int, int]:
        nonlocal track_inclusion

        assert funcname in graphdata.graph

        def putdbg(msg: str) -> None:
            print(f"//dbg-nstatic: {'- '*len(chain)}{msg}")

        node = graphdata.graph[funcname]
        if dbg_nstatic:
            putdbg(f"{funcname}\t{node.nstatic}")
        if node.usage_kind == "dynamic" or node.ndynamic > 0:
            dynamic.add(funcname)
        if track_inclusion:
            included_funcs.add(funcname)

        max_call_nstatic = 0
        max_call_nchain = 0

        if node.calls:
            skipmodel = skipmodels.get(funcname.base())
            chain.append(funcname)
            if len(chain) == cfg_max_call_depth:
                raise ValueError(f"max call depth exceeded: {chain}")
            for call_orig_qname, call_missing_ok in node.calls.items():
                skip_nchain = 0
                # 1. Resolve
                call_qname = graphdata.resolve_funcname(call_orig_qname)
                if not call_qname:
                    if skipmodel:
                        skip, _ = skipmodel(chain[:-1], node, call_orig_qname)
                        if skip:
                            if dbg_nstatic:
                                putdbg(f"{call_orig_qname}\tskip missing")
                            continue
                    if not call_missing_ok:
                        missing.add(call_orig_qname)
                    if dbg_nstatic:
                        putdbg(f"{call_orig_qname}\tmissing")
                    continue

                # 2. Skip
                if skipmodel:
                    skip, skip_nchain = skipmodel(chain[:-1], node, call_qname)
                    max_call_nchain = max(max_call_nchain, skip_nchain)
                    if skip:
                        if dbg_nstatic:
                            putdbg(f"{call_qname}\tskip")
                        continue

                # 3. Call
                if (
                    (not dbg_nocache)
                    and skip_nchain == 0
                    and call_qname in _nstatic_cache
                ):
                    call_nstatic = _nstatic_cache[call_qname]
                    if dbg_nstatic:
                        putdbg(f"{call_qname}\ttotal={call_nstatic} (cache-read)")
                    max_call_nstatic = max(max_call_nstatic, call_nstatic)
                else:
                    call_nstatic, call_nchain = _nstatic(chain, call_qname)
                    max_call_nstatic = max(max_call_nstatic, call_nstatic)
                    max_call_nchain = max(max_call_nchain, call_nchain)
                    if skip_nchain == 0 and call_nchain == 0:
                        if dbg_nstatic:
                            putdbg(f"{call_qname}\ttotal={call_nstatic} (cache-write)")
                        if call_qname not in _nstatic_cache:
                            if dbg_cache:
                                print(f"//dbg-cache: {call_qname} = {call_nstatic}")
                            _nstatic_cache[call_qname] = call_nstatic
                        else:
                            assert dbg_nocache
                            assert _nstatic_cache[call_qname] == call_nstatic
                    elif dbg_nstatic:
                        putdbg(f"{call_qname}\ttotal={call_nstatic} (do-not-cache)")
            chain.pop()
        return node.nstatic + max_call_nstatic, max(0, max_call_nchain - 1)

    def nstatic(funcname: QName) -> int:
        return _nstatic([], funcname)[0]

    groups: dict[str, AnalyzeResultGroup] = {}
    for grp_name, grp_filter in app_func_filters.items():
        rows: dict[QName, AnalyzeResultVal] = {}
        for funcname in graphdata.graph:
            cnt, track_inclusion = grp_filter(funcname)
            if cnt:
                rows[funcname] = AnalyzeResultVal(nstatic=nstatic(funcname), cnt=cnt)
        groups[grp_name] = AnalyzeResultGroup(rows=rows)

    return AnalyzeResult(
        groups=groups, missing=missing, dynamic=dynamic, included_funcs=included_funcs
    )