add callgraph (#103)

this add an analysis constructs callgraph

Author: Roger-luo

src/kirin/analysis/callgraph.py

@@ -0,0 +1,47 @@
+from dataclasses import dataclass, field
+from typing import Iterable
+
+from kirin import ir
+from kirin.dialects import func
+from kirin.print import Printable
+from kirin.print.printer import Printer
+
+
+@dataclass
+class CallGraph(Printable):
+    defs: dict[str, ir.Method] = field(default_factory=dict)
+    backedges: dict[str, set[str]] = field(default_factory=dict)
+
+    def __init__(self, mt: ir.Method):
+        self.defs = {}
+        self.backedges = {}
+        self.__build(mt)
+
+    def __build(self, mt: ir.Method):
+        self.defs[mt.sym_name] = mt
+        for stmt in mt.callable_region.walk():
+            if isinstance(stmt, func.Invoke):
+                backedges = self.backedges.setdefault(stmt.callee.sym_name, set())
+                backedges.add(mt.sym_name)
+                self.__build(stmt.callee)
+
+    def get_neighbors(self, node: str) -> Iterable[str]:
+        return self.backedges.get(node, ())
+
+    def get_edges(self) -> Iterable[tuple[str, str]]:
+        for node, neighbors in self.backedges.items():
+            for neighbor in neighbors:
+                yield node, neighbor
+
+    def get_nodes(self) -> Iterable[str]:
+        return self.defs.keys()
+
+    def print_impl(self, printer: Printer) -> None:
+        for idx, (caller, callee) in enumerate(self.backedges.items()):
+            printer.plain_print(caller)
+            printer.plain_print(" -> ")
+            printer.print_seq(
+                callee, delim=", ", prefix="[", suffix="]", emit=printer.plain_print
+            )
+            if idx < len(self.backedges) - 1:
+                printer.print_newline()

test/analysis/test_callgraph.py

@@ -0,0 +1,36 @@
+from kirin.analysis.callgraph import CallGraph
+from kirin.prelude import basic_no_opt
+
+
+@basic_no_opt
+def abc(a, b):
+    return a + b
+
+
+@basic_no_opt
+def bcd(a, b):
+    return a - b
+
+
+@basic_no_opt
+def cde(a, b):
+    return abc(a, b) + bcd(a, b)
+
+
+@basic_no_opt
+def defg(a, b):
+    return cde(a, b) + abc(a, b)
+
+
+@basic_no_opt
+def efg(a, b):
+    return defg(a, b) + bcd(a, b)
+
+
+def test_callgraph():
+    graph = CallGraph(efg)
+    graph.print()
+    assert "cde" in graph.get_neighbors("abc")
+    assert "defg" in graph.get_neighbors("abc")
+    assert "cde" in graph.get_neighbors("abc")
+    assert "defg" in graph.get_neighbors("abc")