#!/usr/bin/env python3
# build-aux/stack.c.gen - 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 os.path
import re
import sys
import typing

################################################################################
#
# Parse the "VCG" language
#
# https://www.rw.cdl.uni-saarland.de/people/sander/private/html/gsvcg1.html
#
# The formal syntax is found at
# ftp://ftp.cs.uni-sb.de/pub/graphics/vcg/vcg.tgz `doc/grammar.txt`.


class VCGElem:
    typ: str
    lineno: int
    attrs: dict[str, str]


def parse_vcg(reader: typing.TextIO) -> typing.Iterator[VCGElem]:
    re_beg = re.compile(r"(edge|node):\s*\{\s*")
    _re_tok = r"[a-zA-Z_][a-zA-Z0-9_]*"
    _re_str = r'"(?:[^\"]|\\.)*"'
    re_attr = re.compile(
        "(" + _re_tok + r")\s*:\s*(" + _re_tok + "|" + _re_str + r")\s*"
    )
    re_end = re.compile(r"\}\s*$")
    re_skip = re.compile(r"(graph:\s*\{\s*title\s*:\s*" + _re_str + r"\s*|\})\s*")
    re_esc = re.compile(r"\\.")

    for lineno, line in enumerate(reader):
        pos = 0

        def _raise(msg: str) -> typing.NoReturn:
            nonlocal lineno
            nonlocal line
            nonlocal pos
            e = SyntaxError(msg)
            e.lineno = lineno
            e.offset = pos
            e.text = line
            raise e

        if re_skip.fullmatch(line):
            continue

        elem = VCGElem()
        elem.lineno = lineno

        m = re_beg.match(line, pos=pos)
        if not m:
            _raise("does not look like a VCG line")
        elem.typ = m.group(1)
        pos = m.end()

        elem.attrs = {}
        while True:
            if re_end.match(line, pos=pos):
                break
            m = re_attr.match(line, pos=pos)
            if not m:
                _raise("unexpected character")
            k = m.group(1)
            v = m.group(2)
            if k in elem.attrs:
                _raise(f"duplicate key: {repr(k)}")
            if v.startswith('"'):

                def unesc(esc: re.Match[str]) -> str:
                    match esc.group(0)[1:]:
                        case "n":
                            return "\n"
                        case '"':
                            return '"'
                        case "\\":
                            return "\\"
                        case _:
                            _raise(f"invalid escape code {repr(esc.group(0))}")

                v = re_esc.sub(unesc, v[1:-1])
            elem.attrs[k] = v
            pos = m.end()

        yield elem


################################################################################
# Main analysis

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: str
    # .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[str, bool]


def synthetic_node(
    name: str, nstatic: int, calls: typing.Collection[str] = set()
) -> Node:
    n = Node()

    n.funcname = name

    n.location = "<synthetic>"
    n.usage_kind = "static"
    n.nstatic = nstatic
    n.ndynamic = 0

    n.calls = dict((c, False) for c in calls)

    return n


class AnalyzeResultGroup(typing.NamedTuple):
    rows: dict[str, int]
    nmax: int
    nsum: int


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

    included_funcs: set[str]


class Application(typing.Protocol):
    def extra_nodes(self) -> typing.Collection[Node]: ...
    def location_xform(self, loc: str) -> str: ...
    def indirect_callees(self, elem: VCGElem) -> tuple[list[str], bool]: ...
    def skip_call(self, chain: list[str], funcname: str) -> bool: ...


def analyze(
    *,
    ci_fnames: typing.Collection[str],
    app_func_filters: dict[str, typing.Callable[[str], int]],
    app: Application,
    cfg_max_call_depth: int,
) -> AnalyzeResult:
    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,
    )

    graph: dict[str, Node] = dict()
    qualified: dict[str, set[str]] = dict()

    def handle_elem(elem: 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 = v
                        case "label":
                            if elem.attrs.get("shape", "") != "ellipse":
                                m = re_node_label.fullmatch(v)
                                if not m:
                                    raise ValueError(
                                        f"unexpected label value {repr(v)}"
                                    )
                                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 {repr(v)}")
                            skip = True
                        case _:
                            raise ValueError(f"unknown edge key {repr(k)}")
                if not skip:
                    if node.funcname in graph:
                        raise ValueError(f"duplicate node {repr(node.funcname)}")
                    graph[node.funcname] = node
                    if ":" in node.funcname:
                        _, shortname = node.funcname.rsplit(":", 1)
                        if shortname not in qualified:
                            qualified[shortname] = set()
                        qualified[shortname].add(node.funcname)
            case "edge":
                caller: str | None = None
                callee: str | None = None
                for k, v in elem.attrs.items():
                    match k:
                        case "sourcename":
                            caller = v
                        case "targetname":
                            callee = v
                        case "label":
                            pass
                        case _:
                            raise ValueError(f"unknown edge key {repr(k)}")
                if caller is None or callee is None:
                    raise ValueError(f"incomplete edge: {repr(elem.attrs)}")
                if caller not in graph:
                    raise ValueError(f"unknown caller: {caller}")
                if callee == "__indirect_call":
                    callees, missing_ok = app.indirect_callees(elem)
                    for callee in 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 {repr(elem.typ)}")

    for ci_fname in ci_fnames:
        with open(ci_fname, "r") as fh:
            for elem in parse_vcg(fh):
                handle_elem(elem)

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

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

    dbg = False

    def resolve_funcname(funcname: str) -> str | None:
        # Handle `ld --wrap` functions
        if f"__wrap_{funcname}" in graph:
            return f"__wrap_{funcname}"
        if funcname.startswith("__real_") and funcname[len("__real_") :] in graph:
            funcname = funcname[len("__real_") :]

        # Usual case
        if funcname in graph:
            return funcname

        # Handle `__weak` functions
        if funcname in qualified and len(qualified[funcname]) == 1:
            return sorted(qualified[funcname])[0]

        return None

    def nstatic(
        orig_funcname: str, chain: list[str] = [], missing_ok: bool = False
    ) -> int:
        nonlocal dbg
        funcname = resolve_funcname(orig_funcname)
        if not funcname:
            if app.skip_call(chain, orig_funcname):
                return 0
            if not missing_ok:
                missing.add(orig_funcname)
            return 0
        if app.skip_call(chain, funcname):
            return 0

        if len(chain) == cfg_max_call_depth:
            raise ValueError(f"max call depth exceeded: {chain+[funcname]}")

        node = graph[funcname]
        if dbg:
            print(f"//dbg: {funcname}\t{node.nstatic}")
        if node.usage_kind == "dynamic" or node.ndynamic > 0:
            dynamic.add(app.location_xform(funcname))
        included_funcs.add(funcname)
        return node.nstatic + max(
            [
                0,
                *[
                    nstatic(call, chain + [funcname], missing_ok)
                    for call, missing_ok in node.calls.items()
                ],
            ]
        )

    groups: dict[str, AnalyzeResultGroup] = dict()
    for grp_name, grp_filter in app_func_filters.items():
        nmax = 0
        nsum = 0
        rows: dict[str, int] = {}
        for funcname in graph:
            if cnt := grp_filter(funcname):
                n = nstatic(funcname)
                rows[app.location_xform(funcname)] = n
                if n > nmax:
                    nmax = n
                nsum += cnt * n
        groups[grp_name] = AnalyzeResultGroup(rows=rows, nmax=nmax, nsum=nsum)

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


################################################################################
# Mildly-application-specific code


def read_source(location: str) -> str:
    re_location = re.compile(r"(?P<filename>.+):(?P<row>[0-9]+):(?P<col>[0-9]+)")
    m = re_location.fullmatch(location)
    if not m:
        raise ValueError(f"unexpected label value {repr(location)}")
    filename = m.group("filename")
    row = int(m.group("row")) - 1
    col = int(m.group("col")) - 1
    with open(m.group("filename"), "r") as fh:
        return fh.readlines()[row][col:].rstrip()


def get_zero_or_one(
    pred: typing.Callable[[str], bool], fnames: typing.Collection[str]
) -> str | None:
    count = sum(1 for fname in fnames if pred(fname))
    assert count < 2
    if count:
        return next(fname for fname in fnames if pred(fname))
    return None


re_call_other = re.compile(r"(?P<func>[^(]+)\(.*")


class Plugin(typing.Protocol):
    def is_intrhandler(self, name: str) -> bool: ...
    def extra_nodes(self) -> typing.Collection[Node]: ...
    def indirect_callees(
        self, loc: str, line: str
    ) -> tuple[list[str], bool] | None: ...
    def skip_call(self, chain: list[str], call: str) -> bool: ...


class PluginApplication:
    _location_xform: typing.Callable[[str], str]
    _plugins: list[Plugin]

    def __init__(
        self, location_xform: typing.Callable[[str], str], plugins: list[Plugin]
    ) -> None:
        self._location_xform = location_xform
        self._plugins = plugins

    def extra_nodes(self) -> typing.Collection[Node]:
        ret: list[Node] = []
        for plugin in self._plugins:
            ret.extend(plugin.extra_nodes())
        return ret

    def location_xform(self, loc: str) -> str:
        return self._location_xform(loc)

    def indirect_callees(self, elem: VCGElem) -> tuple[list[str], bool]:
        loc = elem.attrs.get("label", "")
        line = read_source(loc)

        for plugin in self._plugins:
            ret = plugin.indirect_callees(loc, line)
            if ret is not None:
                return ret

        placeholder = "__indirect_call"
        if m := re_call_other.fullmatch(line):
            placeholder += ":" + m.group("func")
        placeholder += " at " + self.location_xform(elem.attrs.get("label", ""))
        return [placeholder], False

    def skip_call(self, chain: list[str], funcname: str) -> bool:
        for plugin in self._plugins:
            if plugin.skip_call(chain, funcname):
                return True
        return False


################################################################################
# Application-specific code


class AppPlugin:
    def is_intrhandler(self, name: str) -> bool:
        return False

    def extra_nodes(self) -> typing.Collection[Node]:
        return []

    def indirect_callees(self, loc: str, line: str) -> tuple[list[str], bool] | None:
        if "/3rd-party/" in loc:
            return None
        if "srv->auth" in line:
            return [], False
        if "srv->rootdir" in line:
            return ["get_root"], False
        return None

    def skip_call(self, chain: list[str], call: str) -> bool:
        return False


class LibObjPlugin:
    objcalls: dict[str, set[str]]  # method_name => {method_impls}

    def __init__(self, arg_c_fnames: typing.Collection[str]) -> None:
        ifaces: dict[str, set[str]] = {}  # iface_name => {method_names}
        re_comment = re.compile(r"/\*.*?\*/")
        re_ws = re.compile(r"\s+")
        re_lo_iface = re.compile(r"^\s*#\s*define\s+(?P<name>\S+)_LO_IFACE")
        re_lo_func = re.compile(r"LO_FUNC *\([^,]*, *(?P<name>[^,) ]+) *[,)]")
        for fname in arg_c_fnames:
            with open(fname, "r") as fh:
                while line := fh.readline():
                    if m := re_lo_iface.match(line):
                        iface_name = m.group("name")
                        if iface_name not in ifaces:
                            ifaces[iface_name] = set()
                        while line.endswith("\\\n"):
                            line += fh.readline()
                        line = line.replace("\\\n", " ")
                        line = re_comment.sub(" ", line)
                        line = re_ws.sub(" ", line)
                        for m2 in re_lo_func.finditer(line):
                            ifaces[iface_name].add(m2.group("name"))

        implementations: dict[str, set[str]] = {}  # iface_name => {impl_names}
        for iface_name in ifaces:
            implementations[iface_name] = set()
        re_lo_implementation = re.compile(
            r"^LO_IMPLEMENTATION_[HC]\s*\(\s*(?P<iface>[^, ]+)\s*,\s*(?P<impl_typ>[^,]+)\s*,\s*(?P<impl_name>[^, ]+)\s*[,)].*"
        )
        for fname in arg_c_fnames:
            with open(fname, "r") as fh:
                for line in fh:
                    line = line.strip()
                    if m := re_lo_implementation.match(line):
                        implementations[m.group("iface")].add(m.group("impl_name"))

        objcalls: dict[str, set[str]] = {}  # method_name => {method_impls}
        for iface_name in ifaces:
            for method_name in ifaces[iface_name]:
                if method_name not in objcalls:
                    objcalls[method_name] = set()
                for impl_name in implementations[iface_name]:
                    objcalls[method_name].add(impl_name + "_" + method_name)
        self.objcalls = objcalls

    def is_intrhandler(self, name: str) -> bool:
        return False

    def extra_nodes(self) -> typing.Collection[Node]:
        return []

    def indirect_callees(self, loc: str, line: str) -> tuple[list[str], bool] | None:
        re_call_objcall = re.compile(r"LO_CALL\((?P<obj>[^,]+), (?P<meth>[^,)]+)[,)].*")

        if "/3rd-party/" in loc:
            return None
        if m := re_call_objcall.fullmatch(line):
            if m.group("meth") in self.objcalls:
                return sorted(self.objcalls[m.group("meth")]), False
            return [
                f"__indirect_call:{m.group('obj')}.vtable->{m.group('meth')}"
            ], False
        return None

    def skip_call(self, chain: list[str], call: str) -> bool:
        return False


class LibHWPlugin:
    pico_platform: str

    def __init__(self, arg_pico_platform: str) -> None:
        self.pico_platform = arg_pico_platform

    def is_intrhandler(self, name: str) -> bool:
        return name in [
            "rp2040_hwtimer_intrhandler",
            "hostclock_handle_sig_alarm",
            "hostnet_handle_sig_io",
            "gpioirq_handler",
        ]

    def extra_nodes(self) -> typing.Collection[Node]:
        return []

    def indirect_callees(self, loc: str, line: str) -> tuple[list[str], bool] | None:
        if "/3rd-party/" in loc:
            return None
        if "trigger->cb(trigger->cb_arg)" in line:
            ret = [
                "alarmclock_sleep_intrhandler",
            ]
            if self.pico_platform == "rp2040":
                ret += [
                    "w5500_tcp_alarm_handler",
                    "w5500_udp_alarm_handler",
                ]
            return ret, False
        if "/rp2040_gpioirq.c:" in loc and "handler->fn" in line:
            return [
                "w5500_intrhandler",
            ], False
        return None

    def skip_call(self, chain: list[str], call: str) -> bool:
        return False


class LibCRPlugin:
    def is_intrhandler(self, name: str) -> bool:
        return name in [
            "_cr_gdb_intrhandler",
        ]

    def extra_nodes(self) -> typing.Collection[Node]:
        return []

    def indirect_callees(self, loc: str, line: str) -> tuple[list[str], bool] | None:
        return None

    def skip_call(self, chain: list[str], call: str) -> bool:
        return False


class LibCRIPCPlugin:
    def is_intrhandler(self, name: str) -> bool:
        return False

    def extra_nodes(self) -> typing.Collection[Node]:
        return []

    def indirect_callees(self, loc: str, line: str) -> tuple[list[str], bool] | None:
        if "/3rd-party/" in loc:
            return None
        if "/chan.c:" in loc and "front->dequeue(" in line:
            return [
                "_cr_chan_dequeue",
                "_cr_select_dequeue",
            ], False
        return None

    def skip_call(self, chain: list[str], call: str) -> bool:
        return False


class Lib9PPlugin:
    tmessage_handlers: set[str] | None
    lib9p_msgs: set[str]
    _CONFIG_9P_NUM_SOCKS: int | None
    CONFIG_9P_SRV_MAX_REQS: int | None
    CONFIG_9P_SRV_MAX_DEPTH: int | None

    def __init__(self, arg_base_dir: str, arg_c_fnames: typing.Collection[str]) -> None:
        # Find filenames #######################################################

        def _is_config_h(fname: str) -> bool:
            if not fname.startswith(arg_base_dir + "/"):
                return False
            suffix = fname[len(arg_base_dir) + 1 :]
            if suffix.startswith("3rd-party/"):
                return False
            return suffix.endswith("/config.h")

        config_h_fname = get_zero_or_one(_is_config_h, arg_c_fnames)

        lib9p_srv_c_fname = get_zero_or_one(
            lambda fname: fname.endswith("lib9p/srv.c"), arg_c_fnames
        )

        lib9p_generated_c_fname = get_zero_or_one(
            lambda fname: fname.endswith("lib9p/9p.generated.c"), arg_c_fnames
        )

        # Read config ##########################################################

        def config_h_get(varname: str) -> int | None:
            if config_h_fname:
                with open(config_h_fname, "r") as fh:
                    for line in fh:
                        line = line.rstrip()
                        if line.startswith("#define"):
                            parts = line.split()
                            if parts[1] == varname:
                                return int(parts[2])
            return None

        self._CONFIG_9P_NUM_SOCKS = config_h_get("_CONFIG_9P_NUM_SOCKS")
        self.CONFIG_9P_SRV_MAX_REQS = config_h_get("CONFIG_9P_SRV_MAX_REQS")
        self.CONFIG_9P_SRV_MAX_DEPTH = config_h_get("CONFIG_9P_SRV_MAX_DEPTH")

        # Read sources #########################################################

        tmessage_handlers: set[str] | None = None
        if lib9p_srv_c_fname:
            re_tmessage_handler = re.compile(
                r"^\s*\[LIB9P_TYP_T[^]]+\]\s*=\s*\(tmessage_handler\)\s*(?P<handler>\S+),\s*$"
            )
            tmessage_handlers = set()
            with open(lib9p_srv_c_fname, "r") as fh:
                for line in fh:
                    line = line.rstrip()
                    if m := re_tmessage_handler.fullmatch(line):
                        tmessage_handlers.add(m.group("handler"))
        self.tmessage_handlers = tmessage_handlers

        lib9p_msgs: set[str] = set()
        if lib9p_generated_c_fname:
            re_lib9p_msg_entry = re.compile(r"^\s*_MSG_(?:[A-Z]+)\((?P<typ>\S+)\),$")
            with open(lib9p_generated_c_fname, "r") as fh:
                for line in fh:
                    line = line.rstrip()
                    if m := re_lib9p_msg_entry.fullmatch(line):
                        typ = m.group("typ")
                        lib9p_msgs.add(typ)
        self.lib9p_msgs = lib9p_msgs

    def thread_count(self, name: str) -> int:
        assert self._CONFIG_9P_NUM_SOCKS
        assert self.CONFIG_9P_SRV_MAX_REQS
        if "read" in name:
            return self._CONFIG_9P_NUM_SOCKS
        elif "write" in name:
            return self._CONFIG_9P_NUM_SOCKS * self.CONFIG_9P_SRV_MAX_REQS
        return 1

    def is_intrhandler(self, name: str) -> bool:
        return False

    def extra_nodes(self) -> typing.Collection[Node]:
        return []

    def indirect_callees(self, loc: str, line: str) -> tuple[list[str], bool] | None:
        if "/3rd-party/" in loc:
            return None
        if (
            self.tmessage_handlers
            and "/srv.c:" in loc
            and "tmessage_handlers[typ](" in line
        ):
            # Functions for disabled protocol extensions will be missing.
            return sorted(self.tmessage_handlers), True
        if self.lib9p_msgs and "/9p.c:" in loc:
            for meth in ["validate", "unmarshal", "marshal"]:
                if line.startswith(f"tentry.{meth}("):
                    # Functions for disabled protocol extensions will be missing.
                    return sorted(f"{meth}_{msg}" for msg in self.lib9p_msgs), True
        return None

    def skip_call(self, chain: list[str], call: str) -> bool:
        if "lib9p/srv.c:srv_util_pathfree" in call:
            assert isinstance(self.CONFIG_9P_SRV_MAX_DEPTH, int)
            if len(chain) >= self.CONFIG_9P_SRV_MAX_DEPTH and all(
                ("lib9p/srv.c:srv_util_pathfree" in c)
                for c in chain[-self.CONFIG_9P_SRV_MAX_DEPTH :]
            ):
                return True
        re_msg_meth = re.compile(
            r"^lib9p_(?P<grp>[TR])msg_(?P<meth>validate|unmarshal|marshal)$"
        )
        wrapper = next((c for c in chain if re_msg_meth.match(c)), None)
        if wrapper:
            m = re_msg_meth.match(wrapper)
            assert m
            deny = ":" + m.group("meth") + "_" + ("R" if m.group("grp") == "T" else "T")
            if deny in call:
                return True
        return False


class LibMiscPlugin:
    def is_intrhandler(self, name: str) -> bool:
        return False

    def extra_nodes(self) -> typing.Collection[Node]:
        return []

    def indirect_callees(self, loc: str, line: str) -> tuple[list[str], bool] | None:
        return None

    def skip_call(self, chain: list[str], call: str) -> bool:
        if (
            len(chain) > 1
            and chain[-1] == "__assert_msg_fail"
            and call.endswith(":__lm_printf")
            and "__assert_msg_fail" in chain[:-1]
        ):
            return True
        return False


class PicoSDKPlugin:
    app_init_array: typing.Collection[str]
    app_preinit_array: typing.Collection[str]

    def __init__(
        self,
        *,
        app_init_array: typing.Collection[str],
    ) -> None:
        self.app_init_array = app_init_array

        # git grep '^PICO_RUNTIME_INIT_FUNC\S*('
        self.app_preinit_array = [
            # "runtime_init_mutex",  # pico_mutex
            # "runtime_init_default_alarm_pool",  # pico_time
            # "runtime_init_boot_locks_reset",  # hardware_boot_lock
            "runtime_init_per_core_irq_priorities",  # hardware_irq
            # "spinlock_set_extexclall",  # hardware_sync_spin_lock
            "__aeabi_bits_init",  # pico_bit_ops
            # "runtime_init_bootrom_locking_enable",  # pico_bootrom, rp2350-only
            # "runtime_init_pre_core_tls_setup", # pico_clib_interface, picolibc-only
            # "__aeabi_double_init",  # pico_double
            # "__aeabi_float_init",  # pico_float
            "__aeabi_mem_init",  # pico_mem_ops
            "first_per_core_initializer",  # pico_runtime
            # pico_runtime_init
            # "runtime_init_bootrom_reset",  # rp2350-only
            # "runtime_init_per_core_bootrom_reset",  # rp2350-only
            # "runtime_init_per_core_h3_irq_registers",  # rp2350-only
            "runtime_init_early_resets",
            "runtime_init_usb_power_down",
            # "runtime_init_per_core_enable_coprocessors", # PICO_RUNTIME_SKIP_INIT_PER_CORE_ENABLE_COPROCESSORS
            "runtime_init_clocks",
            "runtime_init_post_clock_resets",
            "runtime_init_rp2040_gpio_ie_disable",
            "runtime_init_spin_locks_reset",
            "runtime_init_install_ram_vector_table",
        ]

    def is_intrhandler(self, name: str) -> bool:
        return name in [
            "isr_invalid",
            "isr_nmi",
            "isr_hardfault",
            "isr_svcall",
            "isr_pendsv",
            "isr_systick",
            *[f"isr_irq{n}" for n in range(32)],
        ]

    def indirect_callees(self, loc: str, line: str) -> tuple[list[str], bool] | None:
        if "/3rd-party/pico-sdk/" not in loc or "/3rd-party/pico-sdk/lib/" in loc:
            return None
        m = re_call_other.fullmatch(line)
        call: str | None = m.group("func") if m else None

        match call:
            case "connect_internal_flash_func":
                return ["rom_func_lookup(ROM_FUNC_CONNECT_INTERNAL_FLASH)"], False
            case "flash_exit_xip_func":
                return ["rom_func_lookup(ROM_FUNC_FLASH_EXIT_XIP)"], False
            case "flash_range_erase_func":
                return ["rom_func_lookup(ROM_FUNC_FLASH_RANGE_ERASE)"], False
            case "flash_flush_cache_func":
                return ["rom_func_lookup(ROM_FUNC_FLASH_FLUSH_CACHE)"], False
            case "rom_table_lookup":
                return ["rom_hword_as_ptr(BOOTROM_TABLE_LOOKUP_OFFSET)"], False
        if "/flash.c:" in loc and "boot2_copyout" in line:
            return ["_stage2_boot"], False
        if "/printf.c:" in loc:
            if call == "out":
                return [
                    "_out_buffer",
                    "_out_null",
                    "_out_fct",
                ], False
            if "->fct(" in line:
                return ["stdio_buffered_printer"], False
        if "/stdio.c:" in loc:
            if call == "out_func":
                return [
                    "stdio_out_chars_crlf",
                    "stdio_out_chars_no_crlf",
                ], False
            if call and (call.startswith("d->") or call.startswith("driver->")):
                _, meth = call.split("->", 1)
                match meth:
                    case "out_chars":
                        return ["stdio_uart_out_chars"], False
                    case "out_flush":
                        return ["stdio_uart_out_flush"], False
                    case "in_chars":
                        return ["stdio_uart_in_chars"], False
        if "/newlib_interface.c:" in loc:
            if line == "*p)();":
                return sorted(self.app_init_array), False
        if "/pico_runtime/runtime.c:" in loc:
            return sorted(self.app_preinit_array), False
        return None

    def skip_call(self, chain: list[str], call: str) -> bool:
        if call == "_out_buffer" or call == "_out_fct":
            last = ""
            for pcall in chain:
                if pcall in [
                    "__wrap_sprintf",
                    "__wrap_snprintf",
                    "__wrap_vsnprintf",
                    "vfctprintf",
                ]:
                    last = pcall
            if last == "vfctprintf":
                return call != "_out_fct"
            else:
                return call == "_out_buffer"
        return False

    def extra_nodes(self) -> typing.Collection[Node]:
        ret = []

        # src/rp2_common/hardware_divider/include/hardware/divider_helper.S
        save_div_state_and_lr = 5 * 4
        # src/rp2_common/pico_divider/divider_hardware.S
        save_div_state_and_lr_64 = 5 * 4

        # src/src/rp2_common/pico_crt0/crt0.S
        for n in range(32):
            ret += [synthetic_node(f"isr_irq{n}", 0, {"__unhandled_user_irq"})]
        ret += [
            synthetic_node("isr_invalid", 0, {"__unhandled_user_irq"}),
            synthetic_node("isr_nmi", 0, {"__unhandled_user_irq"}),
            synthetic_node("isr_hardfault", 0, {"__unhandled_user_irq"}),
            synthetic_node("isr_svcall", 0, {"__unhandled_user_irq"}),
            synthetic_node("isr_pendsv", 0, {"__unhandled_user_irq"}),
            synthetic_node("isr_systick", 0, {"__unhandled_user_irq"}),
            synthetic_node(f"__unhandled_user_irq", 0),
            synthetic_node("_reset_handler", 0, {"runtime_init", "main", "exit"}),
        ]

        ret += [
            # src/rp2_common/pico_int64_ops/pico_int64_ops_aeabi.S
            synthetic_node("__wrap___aeabi_lmul", 4),
            # src/rp2_common/pico_divider/divider_hardware.S
            # s32 aliases
            synthetic_node("div_s32s32", 0, {"divmod_s32s32"}),
            synthetic_node("__wrap___aeabi_idiv", 0, {"divmod_s32s32"}),
            synthetic_node("__wrap___aeabi_idivmod", 0, {"divmod_s32s32"}),
            # s32 impl
            synthetic_node("divmod_s32s32", 0, {"divmod_s32s32_savestate"}),
            synthetic_node(
                "divmod_s32s32_savestate",
                save_div_state_and_lr,
                {"divmod_s32s32_unsafe"},
            ),
            synthetic_node("divmod_s32s32_unsafe", 2 * 4, {"__aeabi_idiv0"}),
            # u32 aliases
            synthetic_node("div_u32u32", 0, {"divmod_u32u32"}),
            synthetic_node("__wrap___aeabi_uidiv", 0, {"divmod_u32u32"}),
            synthetic_node("__wrap___aeabi_uidivmod", 0, {"divmod_u32u32"}),
            # u32 impl
            synthetic_node("divmod_u32u32", 0, {"divmod_u32u32_savestate"}),
            synthetic_node(
                "divmod_u32u32_savestate",
                save_div_state_and_lr,
                {"divmod_u32u32_unsafe"},
            ),
            synthetic_node("divmod_u32u32_unsafe", 2 * 4, {"__aeabi_idiv0"}),
            # s64 aliases
            synthetic_node("div_s64s64", 0, {"divmod_s64s64"}),
            synthetic_node("__wrap___aeabi_ldiv", 0, {"divmod_s64s64"}),
            synthetic_node("__wrap___aeabi_ldivmod", 0, {"divmod_s64s64"}),
            # s64 impl
            synthetic_node("divmod_s64s64", 0, {"divmod_s64s64_savestate"}),
            synthetic_node(
                "divmod_s64s64_savestate",
                save_div_state_and_lr_64 + (2 * 4),
                {"divmod_s64s64_unsafe"},
            ),
            synthetic_node(
                "divmod_s64s64_unsafe", 4, {"divmod_u64u64_unsafe", "__aeabi_ldiv0"}
            ),
            # u64 aliases
            synthetic_node("div_u64u64", 0, {"divmod_u64u64"}),
            synthetic_node("__wrap___aeabi_uldivmod", 0, {"divmod_u64u64"}),
            # u64 impl
            synthetic_node("divmod_u64u64", 0, {"divmod_u64u64_savestate"}),
            synthetic_node(
                "divmod_u64u64_savestate",
                save_div_state_and_lr_64 + (2 * 4),
                {"divmod_u64u64_unsafe"},
            ),
            synthetic_node(
                "divmod_u64u64_unsafe", (1 + 1 + 2 + 5 + 5 + 2) * 4, {"__aeabi_ldiv0"}
            ),
            # *_rem
            synthetic_node("divod_s64s64_rem", 2 * 4, {"divmod_s64s64"}),
            synthetic_node("divod_u64u64_rem", 2 * 4, {"divmod_u64u64"}),
            # src/rp2_common/pico_mem_ops/mem_ops_aeabi.S
            synthetic_node("__aeabi_mem_init", 0, {"rom_funcs_lookup"}),
            synthetic_node(
                "__wrap___aeabi_memset", 0, {"rom_func_lookup(ROM_FUNC_MEMSET)"}
            ),
            synthetic_node("__wrap___aeabi_memset4", 0, {"__wrap___aeabi_memset8"}),
            synthetic_node(
                "__wrap___aeabi_memset8", 0, {"rom_func_lookup(ROM_FUNC_MEMSET4)"}
            ),
            synthetic_node("__wrap___aeabi_memcpy4", 0, {"__wrap___aeabi_memcpy8"}),
            synthetic_node(
                "__wrap___aeabi_memcpy7", 0, {"rom_func_lookup(ROM_FUNC_MEMCPY4)"}
            ),
            synthetic_node("__wrap_memset", 0, {"rom_func_lookup(ROM_FUNC_MEMSET)"}),
            synthetic_node("__wrap___aeabi_memcpy", 0, {"__wrap_memcpy"}),
            synthetic_node("__wrap_memcpy", 0, {"rom_func_lookup(ROM_FUNC_MEMCPY)"}),
            # src/rp2_common/pico_bit_ops/bit_ops_aeabi.S
            synthetic_node("__aeabi_bits_init", 0, {"rom_funcs_lookup"}),
            synthetic_node("__wrap___clz", 0, {"__wrap___clzsi2"}),
            synthetic_node("__wrap___clzl", 0, {"__wrap___clzsi2"}),
            synthetic_node("__wrap___clzsi2", 0, {"rom_func_lookup(ROM_FUNC_CLZ32"}),
            synthetic_node("__wrap___ctzsi2", 0, {"rom_func_lookup(ROM_FUNC_CTZ32"}),
            synthetic_node(
                "__wrap___popcountsi2", 0, {"rom_func_lookup(ROM_FUNC_POPCOUNT32"}
            ),
            synthetic_node("__wrap___clzll", 0, {"__wrap___clzdi2"}),
            synthetic_node("__wrap___clzdi2", 4, {"rom_func_lookup(ROM_FUNC_CLZ32"}),
            synthetic_node("__wrap___ctzdi2", 4, {"rom_func_lookup(ROM_FUNC_CTZ32"}),
            synthetic_node(
                "__wrap___popcountdi2", 3 * 4, {"rom_func_lookup(ROM_FUNC_POPCOUNT32"}
            ),
            synthetic_node("__rev", 0, {"reverse32"}),
            synthetic_node("__revl", 0, {"reverse32"}),
            synthetic_node("reverse32", 0, {"rom_func_lookup(ROM_FUNC_REVERSE32"}),
            synthetic_node("__revll", 0, {"reverse64"}),
            synthetic_node("reverse64", 3 * 4, {"rom_func_lookup(ROM_FUNC_REVERSE32"}),
            # src/rp2040/boot_stage2/boot2_${name,,}.S for name=W25Q080,
            # controlled by `#define PICO_BOOT_STAGE2_{name} 1` in
            # src/boards/include/boards/pico.h
            # synthetic_node("_stage2_boot", 0),  # TODO
            # https://github.com/raspberrypi/pico-bootrom-rp2040
            # synthetic_node("rom_func_lookup(ROM_FUNC_CONNECT_INTERNAL_FLASH)", 0),  # TODO
            # synthetic_node("rom_func_lookup(ROM_FUNC_FLASH_EXIT_XIP)", 0),  # TODO
            # synthetic_node("rom_func_lookup(ROM_FUNC_FLASH_FLUSH_CACHE)", 0),  # TODO
            # synthetic_node("rom_hword_as_ptr(BOOTROM_TABLE_LOOKUP_OFFSET)", 0),  # TODO
        ]
        return ret


class TinyUSBDevicePlugin:
    tud_drivers: dict[str, set[str]]

    def __init__(self, arg_c_fnames: typing.Collection[str]) -> None:
        usbd_c_fname = get_zero_or_one(
            lambda fname: fname.endswith("/tinyusb/src/device/usbd.c"), arg_c_fnames
        )

        tusb_config_h_fname = get_zero_or_one(
            lambda fname: fname.endswith("/tusb_config.h"), arg_c_fnames
        )

        if not usbd_c_fname:
            self.tud_drivers = {}
            return

        assert tusb_config_h_fname
        re_tud_class = re.compile(
            r"^\s*#\s*define\s+(?P<k>CFG_TUD_(?:\S{3}|AUDIO|VIDEO|MIDI|VENDOR|USBTMC|DFU_RUNTIME|ECM_RNDIS))\s+(?P<v>\S+).*"
        )
        tusb_config: dict[str, bool] = {}
        with open(tusb_config_h_fname, "r") as fh:
            in_table = False
            for line in fh:
                line = line.rstrip()
                if m := re_tud_class.fullmatch(line):
                    k = m.group("k")
                    v = m.group("v")
                    tusb_config[k] = bool(int(v))

        tud_drivers: dict[str, set[str]] = {}
        re_tud_entry = re.compile(
            r"^\s+\.(?P<meth>\S+)\s*=\s*(?P<impl>[a-zA-Z0-9_]+)(?:,.*)?"
        )
        re_tud_if1 = re.compile(r"^\s*#\s*if (\S+)\s*")
        re_tud_if2 = re.compile(r"^\s*#\s*if (\S+)\s*\|\|\s*(\S+)\s*")
        re_tud_endif = re.compile(r"^\s*#\s*endif\s*")
        with open(usbd_c_fname, "r") as fh:
            in_table = False
            enabled = True
            for line in fh:
                line = line.rstrip()
                if in_table:
                    if m := re_tud_if1.fullmatch(line):
                        enabled = tusb_config[m.group(1)]
                    elif m := re_tud_if2.fullmatch(line):
                        enabled = tusb_config[m.group(1)] or tusb_config[m.group(2)]
                    elif re_tud_endif.fullmatch(line):
                        enabled = True
                    if m := re_tud_entry.fullmatch(line):
                        meth = m.group("meth")
                        impl = m.group("impl")
                        if meth == "name" or not enabled:
                            continue
                        if meth not in tud_drivers:
                            tud_drivers[meth] = set()
                        if impl != "NULL":
                            tud_drivers[meth].add(impl)
                    if line.startswith("}"):
                        in_table = False
                elif " _usbd_driver[] = {" in line:
                    in_table = True
        self.tud_drivers = tud_drivers

    def is_intrhandler(self, name: str) -> bool:
        return False

    def extra_nodes(self) -> typing.Collection[Node]:
        return []

    def indirect_callees(self, loc: str, line: str) -> tuple[list[str], bool] | None:
        if "/tinyusb/" not in loc or "/tinyusb/src/host/" in loc or "_host.c:" in loc:
            return None
        m = re_call_other.fullmatch(line)
        assert m
        call = m.group("func")
        if call == "_ctrl_xfer.complete_cb":
            return [
                # "process_test_mode_cb",
                "tud_vendor_control_xfer_cb",
                *sorted(self.tud_drivers["control_xfer_cb"]),
            ], False
        elif call.startswith("driver->"):
            return sorted(self.tud_drivers[call[len("driver->") :]]), False
        elif call == "event.func_call.func":
            # callback from usb_defer_func()
            return [], False

        return None

    def skip_call(self, chain: list[str], call: str) -> bool:
        return False


class NewlibPlugin:
    def is_intrhandler(self, name: str) -> bool:
        return False

    def extra_nodes(self) -> typing.Collection[Node]:
        # This is accurate to
        # /usr/arm-none-eabi/lib/thumb/v6-m/nofp/libg.a as of
        # Parabola's arm-none-eabi-newlib 4.5.0.20241231-1.
        return [
            # malloc
            synthetic_node("free", 8, {"_free_r"}),
            synthetic_node("malloc", 8, {"_malloc_r"}),
            synthetic_node("realloc", 8, {"_realloc_r"}),
            synthetic_node("aligned_alloc", 8, {"_memalign_r"}),
            synthetic_node("reallocarray", 24, {"realloc", "__errno"}),
            # synthetic_node("_free_r", 0),  # TODO
            # synthetic_node("_malloc_r", 0),  # TODO
            # synthetic_node("_realloc_r", 0),  # TODO
            # synthetic_node("_memalign_r", 0),  # TODO
            # execution
            synthetic_node("raise", 16, {"_getpid_r"}),
            synthetic_node("abort", 8, {"raise", "_exit"}),
            synthetic_node("longjmp", 0),
            synthetic_node("setjmp", 0),
            # <strings.h>
            synthetic_node("memcmp", 12),
            synthetic_node("memcpy", 28),
            synthetic_node("memset", 20),
            synthetic_node("strcmp", 16),
            synthetic_node("strlen", 8),
            synthetic_node("strncpy", 16),
            synthetic_node("strnlen", 8),
            # other
            synthetic_node("__errno", 0),
            synthetic_node("_getpid_r", 8, {"_getpid"}),
            synthetic_node("random", 8),
            synthetic_node("register_fini", 8, {"atexit"}),
        ]

    def indirect_callees(self, loc: str, line: str) -> tuple[list[str], bool] | None:
        return None

    def skip_call(self, chain: list[str], call: str) -> bool:
        return False


class LibGCCPlugin:
    def is_intrhandler(self, name: str) -> bool:
        return False

    def extra_nodes(self) -> typing.Collection[Node]:
        # This is accurate to
        # /usr/lib/gcc/arm-none-eabi/14.2.0/thumb/v6-m/nofp/libgcc.a
        # as of Parabola's arm-none-eabi-gcc 14.2.0-1.
        return [
            synthetic_node("__aeabi_idiv0", 0),
            synthetic_node("__aeabi_ldiv0", 0),
            synthetic_node("__aeabi_llsr", 0),
        ]

    def indirect_callees(self, loc: str, line: str) -> tuple[list[str], bool] | None:
        return None

    def skip_call(self, chain: list[str], call: str) -> bool:
        return False


def main(
    *,
    arg_pico_platform: str,
    arg_base_dir: str,
    arg_ci_fnames: typing.Collection[str],
    arg_c_fnames: typing.Collection[str],
) -> None:

    plugins: list[Plugin] = []

    # sbc-harness ####################################################

    lib9p_plugin = Lib9PPlugin(arg_base_dir, arg_c_fnames)

    def sbc_is_thread(name: str) -> int:
        if name.endswith("_cr") and name != "lib9p_srv_read_cr":
            if "9p" in name:
                return lib9p_plugin.thread_count(name)
            return 1
        if name == ("_reset_handler" if arg_pico_platform == "rp2040" else "main"):
            return 1
        return 0

    plugins += [
        AppPlugin(),
        LibObjPlugin(arg_c_fnames),
        LibHWPlugin(arg_pico_platform),
        LibCRPlugin(),
        LibCRIPCPlugin(),
        lib9p_plugin,
        LibMiscPlugin(),
    ]

    # pico-sdk #######################################################

    if arg_pico_platform == "rp2040":
        plugins += [
            PicoSDKPlugin(
                app_init_array=["register_fini"],
            ),
            TinyUSBDevicePlugin(arg_c_fnames),
            NewlibPlugin(),
            LibGCCPlugin(),
        ]

    # Tie it all together ############################################

    def thread_filter(name: str) -> int:
        return sbc_is_thread(name)

    def intrhandler_filter(name: str) -> int:
        name = name.rsplit(":", 1)[-1]
        for plugin in plugins:
            if plugin.is_intrhandler(name):
                return 1
        return 0

    def misc_filter(name: str) -> int:
        if name.endswith(":__lm_printf") or name == "__assert_msg_fail":
            return 1
        return 0

    def location_xform(loc: str) -> str:
        if not loc.startswith("/"):
            return loc
        parts = loc.split(":", 1)
        parts[0] = "./" + os.path.relpath(parts[0], arg_base_dir)
        return ":".join(parts)

    result = analyze(
        ci_fnames=arg_ci_fnames,
        app_func_filters={
            "Threads": thread_filter,
            "Interrupt handlers": intrhandler_filter,
            "Misc": misc_filter,
        },
        app=PluginApplication(location_xform, plugins),
        cfg_max_call_depth=100,
    )

    def print_group(grp_name: str) -> None:
        grp = result.groups[grp_name]

        # Figure sizes.
        namelen = max([len(k) for k in grp.rows.keys()] + [len(grp_name) + 4])
        numlen = len(str(grp.nsum))
        sep1 = ("=" * namelen) + " " + "=" * numlen
        sep2 = ("-" * namelen) + " " + "-" * numlen

        # Print.
        print("= " + grp_name + " " + sep1[len(grp_name) + 3 :])
        for name, num in sorted(grp.rows.items()):
            if num == 0:
                continue
            print(f"{name.ljust(namelen)} {str(num).rjust(numlen)}")
        print(sep2)
        print(f"{'Total'.ljust(namelen)} {str(grp.nsum).rjust(numlen)}")
        print(f"{'Maximum'.ljust(namelen)} {str(grp.nmax).rjust(numlen)}")
        print(sep1)

    def next_power_of_2(x: int) -> int:
        return 1 << (x.bit_length())

    print("#include <stddef.h> /* for size_t */")
    print()
    print("/*")
    print_group("Threads")
    print_group("Interrupt handlers")
    print("*/")
    overhead = result.groups["Interrupt handlers"].nmax
    rows: list[tuple[str, int, int]] = []
    for funcname, base in result.groups["Threads"].rows.items():
        rows.append((funcname.split(":")[-1], base, next_power_of_2(base + overhead)))
    namelen = max(len(r[0]) for r in rows)
    baselen = max(len(str(r[1])) for r in rows)
    sizelen = max(len(str(r[2])) for r in rows)
    for row in sorted(rows):
        if row[0] in ("main", "_reset_handler"):
            continue
        print("const size_t CONFIG_COROUTINE_STACK_SIZE_", end="")
        print(f"{row[0].ljust(namelen)} =", end="")
        print(f" {str(row[2]).rjust(sizelen)};", end="")
        print(f" /* LM_NEXT_POWER_OF_2({str(row[1]).rjust(baselen)}+{overhead}) */")
    print()
    print("/*")
    print_group("Misc")

    for funcname in sorted(result.missing):
        print(f"warning: missing: {funcname}")
    for funcname in sorted(result.dynamic):
        print(f"warning: dynamic-stack-usage: {funcname}")

    print("*/")
    print("")
    print("/*")
    for funcname in sorted(result.included_funcs):
        print(f"included: {funcname}")
    print("*/")


if __name__ == "__main__":

    def _main() -> None:
        pico_platform = sys.argv[1]
        base_dir = sys.argv[2]
        obj_fnames = set(sys.argv[3:])

        re_c_obj_suffix = re.compile(r"\.c\.(?:o|obj)$")

        c_fnames: set[str] = set()
        ci_fnames: set[str] = set()
        for obj_fname in obj_fnames:
            if re_c_obj_suffix.search(obj_fname):
                ci_fnames.add(re_c_obj_suffix.sub(".c.ci", obj_fname))
                with open(obj_fname + ".d", "r") as fh:
                    c_fnames.update(
                        fh.read().replace("\\\n", " ").split(":")[-1].split()
                    )

        main(
            arg_pico_platform=pico_platform,
            arg_base_dir=base_dir,
            arg_ci_fnames=ci_fnames,
            arg_c_fnames=c_fnames,
        )

    _main()