feat(cap_graph): Merge similar nodes

jubnzv · jubnzv · commit 1f10a2fe48b2 · 2026-02-01T19:42:14.000+07:00
diff --git a/src/analysis/cap_graph.py b/src/analysis/cap_graph.py
@@ -42,6 +42,89 @@ class CapEdge:
     details: Optional[str] = None  # Additional info
 
 
+def _compute_edge_signature(
+    node_name: str, edges: List[CapEdge]
+) -> Tuple[frozenset, frozenset]:
+    """Compute edge signature for a node (outgoing edges, incoming edges)."""
+    outgoing = frozenset(
+        (e.kind, e.target, e.guard, e.address_class, e.details)
+        for e in edges
+        if e.source == node_name
+    )
+    incoming = frozenset(
+        (e.kind, e.source, e.guard, e.address_class, e.details)
+        for e in edges
+        if e.target == node_name
+    )
+    return (outgoing, incoming)
+
+
+def _merge_equivalent_functions(
+    funcs: List[CapNode], edges: List[CapEdge]
+) -> Tuple[List[CapNode], List[CapEdge], Dict[str, str]]:
+    """
+    Merge function nodes with identical edge signatures.
+
+    Returns:
+        Tuple of (merged_nodes, updated_edges, name_mapping)
+        where name_mapping maps original names to merged node names.
+    """
+    if not funcs:
+        return [], [], {}
+
+    sig_to_funcs: Dict[Tuple[frozenset, frozenset], List[CapNode]] = {}
+    for func in funcs:
+        sig = _compute_edge_signature(func.name, edges)
+        if sig not in sig_to_funcs:
+            sig_to_funcs[sig] = []
+        sig_to_funcs[sig].append(func)
+
+    merged_nodes: List[CapNode] = []
+    name_mapping: Dict[str, str] = {}
+
+    for sig, group in sig_to_funcs.items():
+        if len(group) == 1:
+            merged_nodes.append(group[0])
+            name_mapping[group[0].name] = group[0].name
+        else:
+            sorted_group = sorted(group, key=lambda n: n.name)
+            merged_name = sorted_group[0].name
+            combined_props: Dict[str, Any] = {}
+            for node in sorted_group:
+                combined_props.update(node.properties)
+            combined_props["_merged_names"] = [n.name for n in sorted_group]
+
+            merged_node = CapNode(
+                kind="function",
+                name=merged_name,
+                properties=combined_props,
+            )
+            merged_nodes.append(merged_node)
+            for node in group:
+                name_mapping[node.name] = merged_name
+
+    updated_edges: List[CapEdge] = []
+    seen: Set[Tuple] = set()
+    for edge in edges:
+        new_source = name_mapping.get(edge.source, edge.source)
+        new_target = name_mapping.get(edge.target, edge.target)
+        edge_key = (new_source, new_target, edge.kind, edge.guard, edge.address_class, edge.details)
+        if edge_key not in seen:
+            seen.add(edge_key)
+            updated_edges.append(
+                CapEdge(
+                    kind=edge.kind,
+                    source=new_source,
+                    target=new_target,
+                    guard=edge.guard,
+                    address_class=edge.address_class,
+                    details=edge.details,
+                )
+            )
+
+    return merged_nodes, updated_edges, name_mapping
+
+
 def build_cap_graph(ctx: "ProjectContext") -> Tuple[List[CapNode], List[CapEdge]]:
     """
     Build capability graph from facts.
@@ -301,15 +384,23 @@ def _mermaid_node_id(name: str) -> str:
 def _render_mermaid_node(node: CapNode) -> str:
     """Render a node in Mermaid syntax with appropriate shape."""
     node_id = _mermaid_node_id(node.name)
-    label = node.name.replace('"', r"\"")
 
     if node.kind == "capability":
+        label = node.name.replace('"', r"\"")
         return f'        {node_id}[["{label}"]]'
     elif node.kind == "object":
+        label = node.name.replace('"', r"\"")
         return f'        {node_id}(["{label}"])'
     elif node.kind == "address":
+        label = node.name.replace('"', r"\"")
         return f'        {node_id}{{{{"{label}"}}}}'
     else:  # function
+        merged_names = node.properties.get("_merged_names")
+        if merged_names:
+            func_labels = [f"ƒ {get_simple_name(n)}" for n in merged_names]
+            label = "<br/>".join(func_labels)
+        else:
+            label = f"ƒ {get_simple_name(node.name)}"
         return f'        {node_id}["{label}"]'
 
 
@@ -392,6 +483,10 @@ def dump_cap_graph_to_dir(ctx: "ProjectContext", output_dir: str) -> None:
     funcs = [n for n in nodes if n.kind == "function" and not n.properties.get("init")]
     addrs = [n for n in nodes if n.kind == "address"]
 
+    # Merge functions with identical edge signatures
+    funcs, edges, _ = _merge_equivalent_functions(funcs, edges)
+    init_funcs, edges, _ = _merge_equivalent_functions(init_funcs, edges)
+
     output_path = os.path.join(output_dir, "capgraph.mmd")
     lines = ["flowchart LR"]
 
diff --git a/test/test_cap_graph_merge.py b/test/test_cap_graph_merge.py
@@ -0,0 +1,168 @@
+from analysis.cap_graph import (
+    CapNode,
+    CapEdge,
+    _compute_edge_signature,
+    _merge_equivalent_functions,
+    _render_mermaid_node,
+)
+
+
+class TestCapGraphMerge:
+
+    def test_identical_edges_merged(self):
+        """Functions with identical edge patterns merge into one node."""
+        funcs = [
+            CapNode("function", "mod::foo", {"entry": True}),
+            CapNode("function", "mod::bar", {"entry": True}),
+        ]
+        edges = [
+            CapEdge("mutates", "mod::foo", "SharedObj", guard="role:Admin"),
+            CapEdge("mutates", "mod::bar", "SharedObj", guard="role:Admin"),
+            CapEdge("requires_proof", "mod::foo", "AdminCap"),
+            CapEdge("requires_proof", "mod::bar", "AdminCap"),
+        ]
+
+        merged_funcs, merged_edges, mapping = _merge_equivalent_functions(funcs, edges)
+
+        assert len(merged_funcs) == 1
+        assert "_merged_names" in merged_funcs[0].properties
+        assert set(merged_funcs[0].properties["_merged_names"]) == {"mod::foo", "mod::bar"}
+        assert len(merged_edges) == 2
+
+    def test_different_edges_not_merged(self):
+        """Functions with different edge patterns stay separate."""
+        funcs = [
+            CapNode("function", "mod::foo", {}),
+            CapNode("function", "mod::bar", {}),
+        ]
+        edges = [
+            CapEdge("mutates", "mod::foo", "ObjA"),
+            CapEdge("mutates", "mod::bar", "ObjB"),
+        ]
+
+        merged_funcs, merged_edges, mapping = _merge_equivalent_functions(funcs, edges)
+
+        assert len(merged_funcs) == 2
+        assert mapping["mod::foo"] == "mod::foo"
+        assert mapping["mod::bar"] == "mod::bar"
+
+    def test_partial_overlap_not_merged(self):
+        """Functions with partially overlapping edges stay separate."""
+        funcs = [
+            CapNode("function", "mod::foo", {}),
+            CapNode("function", "mod::bar", {}),
+        ]
+        edges = [
+            CapEdge("mutates", "mod::foo", "SharedObj"),
+            CapEdge("mutates", "mod::bar", "SharedObj"),
+            CapEdge("requires_proof", "mod::foo", "AdminCap"),
+        ]
+
+        merged_funcs, merged_edges, mapping = _merge_equivalent_functions(funcs, edges)
+
+        assert len(merged_funcs) == 2
+
+    def test_incoming_edges_differ_not_merged(self):
+        """Functions with same outgoing but different incoming edges stay separate."""
+        funcs = [
+            CapNode("function", "mod::foo", {}),
+            CapNode("function", "mod::bar", {}),
+        ]
+        edges = [
+            CapEdge("mutates", "mod::foo", "SharedObj"),
+            CapEdge("mutates", "mod::bar", "SharedObj"),
+            CapEdge("calls", "mod::entry", "mod::foo"),
+        ]
+
+        merged_funcs, merged_edges, mapping = _merge_equivalent_functions(funcs, edges)
+
+        assert len(merged_funcs) == 2
+
+    def test_three_funcs_two_merge(self):
+        """Three functions, two with identical edges merge, one stays separate."""
+        funcs = [
+            CapNode("function", "mod::a", {}),
+            CapNode("function", "mod::b", {}),
+            CapNode("function", "mod::c", {}),
+        ]
+        edges = [
+            CapEdge("mutates", "mod::a", "Obj"),
+            CapEdge("mutates", "mod::b", "Obj"),
+            CapEdge("mutates", "mod::c", "OtherObj"),
+        ]
+
+        merged_funcs, merged_edges, mapping = _merge_equivalent_functions(funcs, edges)
+
+        assert len(merged_funcs) == 2
+        merged_names = [f.properties.get("_merged_names") for f in merged_funcs if "_merged_names" in f.properties]
+        assert len(merged_names) == 1
+        assert set(merged_names[0]) == {"mod::a", "mod::b"}
+
+    def test_edge_remapping(self):
+        """Edges get remapped to merged node names."""
+        funcs = [
+            CapNode("function", "mod::foo", {}),
+            CapNode("function", "mod::bar", {}),
+        ]
+        edges = [
+            CapEdge("mutates", "mod::foo", "Obj"),
+            CapEdge("mutates", "mod::bar", "Obj"),
+        ]
+
+        merged_funcs, merged_edges, mapping = _merge_equivalent_functions(funcs, edges)
+
+        assert len(merged_edges) == 1
+        assert merged_edges[0].source == merged_funcs[0].name
+
+    def test_no_funcs_empty_result(self):
+        """Empty input returns empty output."""
+        merged_funcs, merged_edges, mapping = _merge_equivalent_functions([], [])
+
+        assert merged_funcs == []
+        assert merged_edges == []
+        assert mapping == {}
+
+    def test_single_func_no_merge(self):
+        """Single function stays as-is."""
+        funcs = [CapNode("function", "mod::only", {"entry": True})]
+        edges = [CapEdge("mutates", "mod::only", "Obj")]
+
+        merged_funcs, merged_edges, mapping = _merge_equivalent_functions(funcs, edges)
+
+        assert len(merged_funcs) == 1
+        assert "_merged_names" not in merged_funcs[0].properties
+        assert mapping["mod::only"] == "mod::only"
+
+
+class TestMermaidRendering:
+
+    def test_merged_node_renders_with_function_symbol(self):
+        """Merged node renders with ƒ prefix and <br/> separators."""
+        node = CapNode("function", "pkg::mod::foo", {
+            "_merged_names": ["pkg::mod::foo", "pkg::mod::bar", "pkg::mod::baz"]
+        })
+
+        result = _render_mermaid_node(node)
+
+        assert "ƒ foo" in result
+        assert "ƒ bar" in result
+        assert "ƒ baz" in result
+        assert "<br/>" in result
+
+    def test_single_func_renders_with_function_symbol(self):
+        """Single function node renders with ƒ prefix."""
+        node = CapNode("function", "pkg::mod::my_func", {})
+
+        result = _render_mermaid_node(node)
+
+        assert "ƒ my_func" in result
+        assert "<br/>" not in result
+
+    def test_capability_no_function_symbol(self):
+        """Capability nodes don't get ƒ symbol."""
+        node = CapNode("capability", "pkg::mod::AdminCap", {})
+
+        result = _render_mermaid_node(node)
+
+        assert "ƒ" not in result
+        assert "AdminCap" in result
