feat: supprt function walk and cg

Author: SunBK201

scubatrace/file.py

@@ -326,6 +326,18 @@ def function_by_line(self, line: int) -> Function | None:
                 return func
         return None
 
+    def functions_by_name(self, name: str) -> list[Function]:
+        """
+        The functions that have the specified name.
+
+        Args:
+            name (str): The name of the function to check.
+
+        Returns:
+            list[Function]: A list of functions that have the specified name.
+        """
+        return [f for f in self.functions if f.name == name]
+
     def statements_by_line(self, line: int) -> list[Statement]:
         """
         The statements that are located on the specified line number.

scubatrace/function.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
-from collections import defaultdict
+from collections import defaultdict, deque
 from functools import cached_property
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable, Generator
 
 import networkx as nx
 from tree_sitter import Node
@@ -102,7 +102,9 @@ def first_statement(self) -> Statement | None:
         return self.statements[0]
 
     def __str__(self) -> str:
-        return self.signature
+        return (
+            f'"{self.name.replace("::", "--")} ({self.file.name}\\:{self.start_line})"'
+        )
 
     def set_joernid(self, joern_id: str):
         self.joern_id = joern_id
@@ -318,6 +320,106 @@ def callers(self) -> dict[Function, list[Statement]]:
                         break
         return callers
 
+    def walk_backward(
+        self,
+        filter: Callable[[Statement], bool] | None = None,
+        stop_by: Callable[[Statement], bool] | None = None,
+        depth: int = -1,
+        base: str = "call",
+    ) -> Generator[Function, None, None]:
+        for caller in super().walk_backward(
+            filter=filter,
+            stop_by=stop_by,
+            depth=depth,
+            base=base,
+        ):
+            assert isinstance(caller, Function)
+            yield caller
+
+    def walk_forward(
+        self,
+        filter: Callable[[Statement], bool] | None = None,
+        stop_by: Callable[[Statement], bool] | None = None,
+        depth: int = -1,
+        base: str = "call",
+    ) -> Generator[Function, None, None]:
+        for callee in super().walk_forward(
+            filter=filter,
+            stop_by=stop_by,
+            depth=depth,
+            base=base,
+        ):
+            assert isinstance(callee, Function)
+            yield callee
+
+    def __build_callgraph(self, depth: int = -1) -> nx.MultiDiGraph:
+        cg = nx.MultiDiGraph()
+        cg.add_node(
+            self,
+            color="red",
+            shape="box",
+            style="rounded",
+        )
+        forward_depth = 2048 if depth == -1 else depth
+        dq: deque[Function | FunctionDeclaration] = deque([self])
+        visited: set[Function | FunctionDeclaration] = set([self])
+        while len(dq) > 0 and forward_depth > 0:
+            size = len(dq)
+            for _ in range(size):
+                caller = dq.popleft()
+                if not isinstance(caller, Function):
+                    continue
+                for callee, callsites in caller.callees.items():
+                    for callsite in callsites:
+                        cg.add_edge(
+                            caller,
+                            callee,
+                            key=callsite.signature,
+                            label=callsite.start_line,
+                        )
+                    if callee not in visited:
+                        visited.add(callee)
+                        dq.append(callee)
+            forward_depth -= 1
+
+        backward_depth = 2048 if depth == -1 else depth
+        dq = deque([self])
+        visited = set([self])
+        while len(dq) > 0 and backward_depth > 0:
+            size = len(dq)
+            for _ in range(size):
+                callee = dq.popleft()
+                if not isinstance(callee, Function):
+                    continue
+                for caller, callsites in callee.callers.items():
+                    for callsite in callsites:
+                        cg.add_edge(
+                            caller,
+                            callee,
+                            key=callsite.signature,
+                            label=callsite.start_line,
+                        )
+                    if caller not in visited:
+                        visited.add(caller)
+                        dq.append(caller)
+            backward_depth -= 1
+        return cg
+
+    def export_callgraph(self, path: str, depth: int = -1) -> nx.MultiDiGraph:
+        """
+        Exports the call graph of the function to a DOT file.
+
+        Args:
+            path (str): The path to save the DOT file.
+            depth (int): The depth of the call graph to export. -1 means no limit.
+
+        Returns:
+            nx.MultiDiGraph: The call graph of the function.
+        """
+        cg = self.__build_callgraph(depth)
+        nx.nx_pydot.write_dot(cg, path)
+        return cg
+
     def slice_by_statements(
         self,
         statements: list[Statement],

scubatrace/statement.py

@@ -514,6 +514,11 @@ def walk_backward(
                             nexts.extend(stats)
                     case "control_dependent":
                         nexts = cur_stat.pre_control_dependents
+                    case "call":
+                        from .function import Function
+
+                        assert isinstance(cur_stat, Function)
+                        nexts = [caller for caller in cur_stat.callers.keys()]
                     case _:
                         nexts = cur_stat.pre_controls
                 for pre in nexts:
@@ -539,7 +544,7 @@ def walk_forward(
             stop_by (Callable[[Statement], bool] | None): A function to stop the walking when it returns True.
             depth (int): The maximum depth to walk forward. Default is -1, which means no limit.
             base (str): The base type of the walk.
-                Can be "control", "data_dependent", or "control_dependent".
+                Can be "control", "data_dependent", "control_dependent", "call".
 
         Yields:
             Statement: The statements that match the filter or all statements if no filter is provided.
@@ -564,11 +569,18 @@ def walk_forward(
                             nexts.extend(stats)
                     case "control_dependent":
                         nexts = cur_stat.post_control_dependents
+                    case "call":
+                        from .function import Function
+
+                        assert isinstance(cur_stat, Function)
+                        nexts = [caller for caller in cur_stat.callees.keys()]
                     case _:
                         nexts = cur_stat.post_controls
                 for post in nexts:
                     if post in visited:
                         continue
+                    if not isinstance(post, Statement):
+                        continue
                     visited.add(post)
                     dq.appendleft(post)
             depth -= 1

tests/samples/c/main.c

@@ -1,9 +1,11 @@
 #include "include/sub.h"
 #include <stdio.h>
 
+int mul(int a, int b);
+
 int add(int a, int b)
 {
-    return a + sub(a, b);
+    return a + sub(a, b) + mul(a, b);
 }
 
 int main(int argc, char** argv)

tests/test_file.py

@@ -36,26 +36,32 @@ def test_file_functions(self):
             self.assertIsNotNone(func.name)
 
     def test_file_function_by_line(self):
-        function = self.file.function_by_line(5)
+        function = self.file.function_by_line(8)
         self.assertIsNotNone(function)
         assert function is not None
         self.assertEqual(function.name, "add")
 
+    def test_file_function_by_name(self):
+        function = self.file.functions_by_name("add")
+        self.assertEqual(len(function), 1)
+        function = self.file.functions_by_name("main")
+        self.assertEqual(len(function), 1)
+
     def test_file_statements(self):
         statements = self.file.statements
         self.assertGreater(len(statements), 0)
         for stmt in statements:
             self.assertIsNotNone(stmt.text)
 
     def test_file_statement_by_line(self):
-        statements = self.file.statements_by_line(14)
+        statements = self.file.statements_by_line(16)
         self.assertGreater(len(statements), 0)
         self.assertEqual(statements[0].text, "int c = count + argc;")
 
         self.assertEqual(len(self.file.statements_by_line(-1)), 0)
         self.assertGreater(len(self.file.statements_by_line(1)), 0)
 
-        self.assertEqual(self.file.statements_by_line(18)[0].text, "a -= 1;")
+        self.assertEqual(self.file.statements_by_line(20)[0].text, "a -= 1;")
 
     def test_file_identifiers(self):
         identifiers = self.file.identifiers
@@ -71,7 +77,7 @@ def test_file_variables(self):
 
     def test_file_cfg(self):
         assert self.file is not None
-        cfg = self.file.export_cfg_dot(f"{self.file.name}.dot")
+        cfg = self.file.export_cfg_dot(f"{self.project_path}/{self.file.name}.dot")
         self.assertIsNotNone(cfg)
         self.assertGreater(len(cfg.nodes), 0)
         self.assertGreater(len(cfg.edges), 0)
@@ -84,7 +90,7 @@ def test_file_query(self):
             )@call
         """
         query = self.file.query(query_str)
-        target_lines = [6, 36]
+        target_lines = [8, 38]
         self.assertEqual(len(query), len(target_lines))
         for stat in query:
             self.assertIn(stat.start_line, target_lines)

tests/test_function.py

@@ -15,9 +15,7 @@ def setUp(self):
         file = self.project.files.get("main.c")
         assert file is not None
         self.file = file
-        function = self.file.function_by_line(11)
-        assert function is not None
-        self.function = function
+        self.function = self.file.functions_by_name("main")[0]
 
     def test_function_create(self):
         function = scubatrace.Function.create(self.function.node, self.function.parent)
@@ -31,7 +29,7 @@ def test_function_callees(self):
         self.assertIn("printf", [callee.name for callee in callees])
 
     def test_function_callers(self):
-        function = self.file.function_by_line(4)
+        function = self.file.function_by_line(6)
         self.assertIsNotNone(function)
         assert function is not None
         callers = function.callers
@@ -43,7 +41,7 @@ def test_function_lines(self):
 
     def test_function_parameter_lines(self):
         self.assertEqual(len(self.function.parameter_lines), 1)
-        self.assertEqual(self.function.parameter_lines[0], 9)
+        self.assertEqual(self.function.parameter_lines[0], 11)
 
     def test_function_parameters(self):
         parameters = self.function.parameters
@@ -58,10 +56,40 @@ def test_function_variables(self):
         self.assertEqual(variables[len(variables) - 1].text, "count")
 
     def test_function_export_cfg_dot(self):
-        cfg = self.function.export_cfg_dot("cfg.dot", with_cdg=True, with_ddg=True)
+        cfg = self.function.export_cfg_dot(
+            f"{self.project_path}/{self.function.name}.dot",
+            with_cdg=True,
+            with_ddg=True,
+        )
         self.assertIsNotNone(cfg)
 
     def test_function_slicing_by_lines(self):
-        stats = self.function.slice_by_lines([14])
-        self.assertEqual(stats[0].start_line, 9)
-        self.assertEqual(stats[len(stats) - 1].start_line, 38)
+        stats = self.function.slice_by_lines([16])
+        self.assertEqual(stats[0].start_line, 11)
+        self.assertEqual(stats[len(stats) - 1].start_line, second=40)
+
+    def test_function_walk_backward(self):
+        function = self.file.functions_by_name("add")[0]
+        assert function is not None
+        functions = list(function.walk_backward())
+        self.assertEqual(len(functions), 3)
+        functions_start_lines = sorted([f.start_line for f in functions])
+        self.assertEqual(functions_start_lines, [6, 11, 51])
+
+    def test_function_walk_forward(self):
+        functions = list(self.function.walk_forward())
+        self.assertGreater(len(functions), 0)
+        functions_start_lines = sorted([f.start_line for f in functions])
+        self.assertIn(11, functions_start_lines)
+        self.assertIn(6, functions_start_lines)
+        self.assertIn(51, functions_start_lines)
+
+    def test_function_export_callgraph(self):
+        function = self.file.functions_by_name("add")[0]
+        callgraph = function.export_callgraph(
+            f"{self.project_path}/{function.name}_callgraph.dot",
+            depth=1,
+        )
+        self.assertIsNotNone(callgraph)
+        self.assertGreater(len(callgraph.nodes), 0)
+        self.assertGreater(len(callgraph.edges), 0)

tests/test_identifer.py

@@ -15,7 +15,7 @@ def setUp(self):
         self.file = self.project.files.get("main.c")
         assert self.file is not None
         self.assertGreater(len(self.file.statements), 0)
-        self.statement = self.file.statements_by_line(14)[0]
+        self.statement = self.file.statements_by_line(16)[0]
         self.assertGreater(len(self.statement.identifiers), 0)
         self.identifier = self.statement.identifiers[0]
 
@@ -27,4 +27,4 @@ def test_identifier_post_data_dependents(self):
         dependents = self.identifier.post_data_dependents
         self.assertEqual(len(dependents), 5)
         dependents_lines = sorted([dep.start_line for dep in dependents])
-        self.assertEqual(dependents_lines, [15, 15, 34, 36, 38])
+        self.assertEqual(dependents_lines, [17, 17, 36, 38, 40])