feat(dataflow): identity map — IR node ids ↔ ordinal can:// ids

Author: rahlk

codeanalyzer/dataflow/identity.py

@@ -0,0 +1,31 @@
+"""Bijection between internal IR node ids (ints, per function) and canonical
+ordinal ids `<callable can:// id>@<tag>` — `@entry`/`@exit` for the synthetic
+CFG bookends, `@line:col` for real statements. Both emitters consume this so
+JSON body-node ids and Neo4j PyCFGNode keys are identical."""
+from __future__ import annotations
+from typing import Dict, Iterable
+
+
+class IdentityMap:
+    def __init__(self, callable_id: str, id_to_ordinal: Dict[int, str]):
+        self._callable_id = callable_id
+        self._map = id_to_ordinal
+
+    @classmethod
+    def for_function(cls, callable_id: str, pdg) -> "IdentityMap":
+        cfg = pdg.cfg
+        m: Dict[int, str] = {}
+        for n in cfg.nodes:
+            if n.id == cfg.entry_id:
+                m[n.id] = f"{callable_id}@entry"
+            elif n.id == cfg.exit_id:
+                m[n.id] = f"{callable_id}@exit"
+            else:
+                m[n.id] = f"{callable_id}@{n.start_line}:{n.start_column}"
+        return cls(callable_id, m)
+
+    def ordinal(self, node_id: int) -> str:
+        return self._map[node_id]
+
+    def node_ids(self) -> Iterable[int]:
+        return self._map.keys()

test/test_v2_l3.py

@@ -0,0 +1,23 @@
+from codeanalyzer.dataflow.identity import IdentityMap
+
+class _Node:
+    def __init__(self, id, start_line, start_column, kind):
+        self.id, self.start_line, self.start_column, self.kind = id, start_line, start_column, kind
+
+class _CFG:
+    def __init__(self, nodes, entry_id, exit_id):
+        self._n = {n.id: n for n in nodes}; self.nodes = nodes
+        self.entry_id, self.exit_id = entry_id, exit_id
+    def node_by_id(self, i): return self._n[i]
+
+class _PDG:
+    def __init__(self, cfg): self.cfg = cfg
+
+def test_ordinal_ids_for_entry_exit_and_statements():
+    nodes = [_Node(0, 1, 0, "entry"), _Node(1, 2, 4, "statement"), _Node(2, 3, 4, "exit")]
+    pdg = _PDG(_CFG(nodes, entry_id=0, exit_id=2))
+    im = IdentityMap.for_function("can://python/app/m.py/f()", pdg)
+    assert im.ordinal(0) == "can://python/app/m.py/f()@entry"
+    assert im.ordinal(1) == "can://python/app/m.py/f()@2:4"
+    assert im.ordinal(2) == "can://python/app/m.py/f()@exit"
+    assert set(im.node_ids()) == {0, 1, 2}