measurestack: Rework the nstatic memoization to be easier to think aboutlukeshu/measurestack-refactor

1 files changed, 120 insertions, 77 deletions

diff --git a/build-aux/measurestack/app_plugins.py b/build-aux/measurestack/app_plugins.py
index e365f82..12c0384 100644
--- a/build-aux/measurestack/app_plugins.py
+++ b/build-aux/measurestack/app_plugins.py
@@ -8,7 +8,7 @@ import subprocess
 import typing
 
 from . import analyze, util
-from .analyze import BaseName, Node, QName
+from .analyze import BaseName, Node, NodeHandleCB, QName
 from .util import synthetic_node
 
 # pylint: disable=unused-variable
@@ -51,7 +51,7 @@ class CmdPlugin:
             return [QName("get_root")], False
         return None
 
-    def skipmodels(self) -> dict[BaseName, analyze.SkipModel]:
+    def node_handlers(self) -> dict[BaseName, analyze.NodeHandler]:
         return {}
 
 
@@ -128,22 +128,25 @@ class LibMiscPlugin:
             ], False
         return None
 
-    def skipmodels(self) -> dict[BaseName, analyze.SkipModel]:
+    def node_handlers(self) -> dict[BaseName, analyze.NodeHandler]:
+        def handle___assert_msg_fail(
+            handle: NodeHandleCB, chain: typing.Sequence[QName], node: Node
+        ) -> tuple[int, bool]:
+            inner: bool = any(c.base() == BaseName("__assert_msg_fail") for c in chain)
+
+            def skip_p(call_qname: QName) -> bool:
+                return inner and call_qname.base() in [
+                    BaseName("__lm_printf"),
+                    BaseName("__lm_light_printf"),
+                ]
+
+            nstatic, _ = analyze.handle_simple_node(handle, chain, node, skip_p)
+            return nstatic, not inner
+
         return {
-            BaseName("__assert_msg_fail"): analyze.SkipModel(
-                {BaseName("__assert_msg_fail")}, self._skipmodel___assert_msg_fail
-            ),
+            BaseName("__assert_msg_fail"): handle___assert_msg_fail,
         }
 
-    def _skipmodel___assert_msg_fail(
-        self, chain: typing.Sequence[QName], node: Node, call: QName
-    ) -> bool:
-        if call.base() in [BaseName("__lm_printf"), BaseName("__lm_light_printf")]:
-            return any(
-                c.base() == BaseName("__assert_msg_fail") for c in reversed(chain)
-            )
-        return False
-
 
 class LibHWPlugin:
     pico_platform: str
@@ -210,7 +213,7 @@ class LibHWPlugin:
             ], False
         return None
 
-    def skipmodels(self) -> dict[BaseName, analyze.SkipModel]:
+    def node_handlers(self) -> dict[BaseName, analyze.NodeHandler]:
         return {}
 
 
@@ -234,7 +237,7 @@ class LibCRPlugin:
     ) -> tuple[typing.Collection[QName], bool] | None:
         return None
 
-    def skipmodels(self) -> dict[BaseName, analyze.SkipModel]:
+    def node_handlers(self) -> dict[BaseName, analyze.NodeHandler]:
         return {}
 
 
@@ -263,7 +266,7 @@ class LibCRIPCPlugin:
             ], False
         return None
 
-    def skipmodels(self) -> dict[BaseName, analyze.SkipModel]:
+    def node_handlers(self) -> dict[BaseName, analyze.NodeHandler]:
         return {}
 
 
@@ -272,12 +275,6 @@ class Lib9PPlugin:
         r"^\s*\[LIB9P_TYP_T[^]]+\]\s*=\s*\(tmessage_handler\)\s*(?P<handler>\S+),\s*$"
     )
     re_lib9p_msg_entry = re.compile(r"^\s*_MSG_(?:[A-Z]+)\((?P<typ>\S+)\),$")
-    re_lib9p_caller = re.compile(
-        r"^lib9p_(?P<grp>[TR])msg_(?P<meth>validate|unmarshal|marshal)$"
-    )
-    re_lib9p_callee = re.compile(
-        r"^(?P<meth>validate|unmarshal|marshal)_(?P<msg>(?P<grp>[TR]).*)$"
-    )
 
     tmessage_handlers: set[QName] | None
     lib9p_msgs: set[str]
@@ -396,33 +393,51 @@ class Lib9PPlugin:
                     return [QName(f"{meth}_{msg}") for msg in self.lib9p_msgs], True
         return None
 
-    def skipmodels(self) -> dict[BaseName, analyze.SkipModel]:
-        ret: dict[BaseName, analyze.SkipModel] = {
-            BaseName("_lib9p_validate"): analyze.SkipModel(
-                1,
-                self._skipmodel__lib9p_validate_unmarshal_marshal,
-            ),
-            BaseName("_lib9p_unmarshal"): analyze.SkipModel(
-                1,
-                self._skipmodel__lib9p_validate_unmarshal_marshal,
-            ),
-            BaseName("_lib9p_marshal"): analyze.SkipModel(
-                1,
-                self._skipmodel__lib9p_validate_unmarshal_marshal,
-            ),
-        }
-        return ret
+    re_lib9p_caller = re.compile(
+        r"^lib9p_(?P<grp>[TR])msg_(?P<meth>validate|unmarshal|marshal)$"
+    )
+    re_lib9p_callee = re.compile(
+        r"^(?P<meth>validate|unmarshal|marshal)_(?P<msg>(?P<grp>[TR]).*)$"
+    )
 
-    def _skipmodel__lib9p_validate_unmarshal_marshal(
-        self, chain: typing.Sequence[QName], node: Node, call: QName
-    ) -> bool:
-        m_caller = self.re_lib9p_caller.fullmatch(str(chain[-1].base()))
-        assert m_caller
+    def node_handlers(self) -> dict[BaseName, analyze.NodeHandler]:
+        def handle_9p_inner(
+            handle: NodeHandleCB, chain: typing.Sequence[QName], node: Node
+        ) -> tuple[int, bool]:
+            m_caller = self.re_lib9p_caller.fullmatch(str(chain[-1].base()))
+            assert m_caller
+            grp = m_caller.group("grp")
+
+            def skip_p(call_qname: QName) -> bool:
+                if m_callee := self.re_lib9p_callee.fullmatch(str(call_qname.base())):
+                    if m_callee.group("grp") != grp:
+                        return True
+                return False
 
-        m_callee = self.re_lib9p_callee.fullmatch(str(call.base()))
-        if not m_callee:
-            return False
-        return m_caller.group("grp") != m_callee.group("grp")
+            nstatic, _ = analyze.handle_simple_node(handle, chain, node, skip_p)
+            return nstatic, False
+
+        def handle_9p_outer(
+            handle: NodeHandleCB, chain: typing.Sequence[QName], node: Node
+        ) -> tuple[int, bool]:
+            nstatic, _ = analyze.handle_simple_node(handle, chain, node)
+            return nstatic, True
+
+        ret: dict[BaseName, analyze.NodeHandler] = {
+            # validate
+            BaseName("lib9p_Tmsg_validate"): handle_9p_outer,
+            BaseName("lib9p_Rmsg_validate"): handle_9p_outer,
+            BaseName("_lib9p_validate"): handle_9p_inner,
+            # unmarsahl
+            BaseName("lib9p_Tmsg_unmarshal"): handle_9p_outer,
+            BaseName("lib9p_Rmsg_unmarshal"): handle_9p_outer,
+            BaseName("_lib9p_unmarshal"): handle_9p_inner,
+            # marshal
+            BaseName("lib9p_Tmsg_marshal"): handle_9p_outer,
+            BaseName("lib9p_Rmsg_marshal"): handle_9p_outer,
+            BaseName("_lib9p_marshal"): handle_9p_inner,
+        }
+        return ret
 
 
 class PicoFmtPlugin:
@@ -495,35 +510,63 @@ class PicoFmtPlugin:
                 ], False
         return None
 
-    def skipmodels(self) -> dict[BaseName, analyze.SkipModel]:
-        ret: dict[BaseName, analyze.SkipModel] = {
-            BaseName("fmt_state_putchar"): analyze.SkipModel(
-                self.known_fct.keys(), self._skipmodel_fmt_state_putchar
-            ),
-            BaseName("_vfctprintf"): analyze.SkipModel(
-                self.wont_call_v, self._skipmodel__vfctprintf
-            ),
-        }
-        return ret
+    def node_handlers(self) -> dict[BaseName, analyze.NodeHandler]:
+        putchar_cache: dict[BaseName | None, int] = {}  # fct=>total_nstatic
 
-    def _skipmodel_fmt_state_putchar(
-        self, chain: typing.Sequence[QName], node: Node, call: QName
-    ) -> bool:
-        if call.base() in self.known_fct.values():
+        def handle_putchar(
+            handle: NodeHandleCB, chain: typing.Sequence[QName], node: Node
+        ) -> tuple[int, bool]:
             fct: BaseName | None = None
-            for pcall in reversed(chain):
+            cachekey: BaseName | None = None
+            for i in reversed(range(len(chain))):
+                pcall = chain[i]
                 if pcall.base() in self.known_fct:
                     fct = self.known_fct[pcall.base()]
-                    return call.base() != fct
-            return True
-        return False
+                    cachekey = fct
+                    if fct == BaseName("libfmt_light_fct"):
+                        assert pcall.base() == BaseName("__lm_light_printf")
+                        if i > 0 and chain[i - 1].base() == BaseName(
+                            "__assert_msg_fail"
+                        ):
+                            cachekey = BaseName("inner/" + str(cachekey))
+                    break
+            if cachekey in putchar_cache:
+                return putchar_cache[cachekey], False
+
+            def skip_p(call_qname: QName) -> bool:
+                return fct is None or call_qname.base() != fct
+
+            nstatic, _ = analyze.handle_simple_node(handle, chain, node, skip_p)
+            putchar_cache[cachekey] = nstatic
+            return nstatic, False
+
+        def handle__vfctprintf(
+            handle: NodeHandleCB, chain: typing.Sequence[QName], node: Node
+        ) -> tuple[int, bool]:
+            def skip_p(call_qname: QName) -> bool:
+                if call_qname.base() == BaseName("libfmt_conv_formatter"):
+                    return any(c.base() in self.wont_call_v for c in chain)
+                return False
 
-    def _skipmodel__vfctprintf(
-        self, chain: typing.Sequence[QName], node: Node, call: QName
-    ) -> bool:
-        if call.base() == BaseName("libfmt_conv_formatter"):
-            return any(c.base() in self.wont_call_v for c in chain)
-        return False
+            nstatic, _ = analyze.handle_simple_node(handle, chain, node, skip_p)
+            return nstatic, False
+
+        def handle_top(
+            handle: NodeHandleCB, chain: typing.Sequence[QName], node: Node
+        ) -> tuple[int, bool]:
+            nstatic, _ = analyze.handle_simple_node(handle, chain, node)
+            in_fail: bool = any(
+                c.base() == BaseName("__assert_msg_fail") for c in chain
+            )
+            return nstatic, not in_fail
+
+        ret: dict[BaseName, analyze.NodeHandler] = {
+            BaseName("fmt_state_putchar"): handle_putchar,
+            BaseName("_vfctprintf"): handle__vfctprintf,
+        }
+        for top in self.known_fct:
+            ret[top] = handle_top
+        return ret
 
 
 class PicoSDKPlugin:
@@ -632,7 +675,7 @@ class PicoSDKPlugin:
             return self.app_preinit_array, False
         return None
 
-    def skipmodels(self) -> dict[BaseName, analyze.SkipModel]:
+    def node_handlers(self) -> dict[BaseName, analyze.NodeHandler]:
         return {}
 
     def extra_nodes(self) -> typing.Collection[Node]:
@@ -865,7 +908,7 @@ class TinyUSBDevicePlugin:
 
         return None
 
-    def skipmodels(self) -> dict[BaseName, analyze.SkipModel]:
+    def node_handlers(self) -> dict[BaseName, analyze.NodeHandler]:
         return {}
 
 
@@ -932,7 +975,7 @@ class NewlibPlugin:
     ) -> tuple[typing.Collection[QName], bool] | None:
         return None
 
-    def skipmodels(self) -> dict[BaseName, analyze.SkipModel]:
+    def node_handlers(self) -> dict[BaseName, analyze.NodeHandler]:
         return {}
 
 
@@ -965,5 +1008,5 @@ class LibGCCPlugin:
     ) -> tuple[typing.Collection[QName], bool] | None:
         return None
 
-    def skipmodels(self) -> dict[BaseName, analyze.SkipModel]:
+    def node_handlers(self) -> dict[BaseName, analyze.NodeHandler]:
         return {}