fix: fix defer control flow in go

SunBK201 · SunBK201 · commit db3493a9ea91 · 2025-08-02T01:49:07.000+08:00
diff --git a/scubatrace/function.py b/scubatrace/function.py
@@ -96,6 +96,12 @@ def statements(self) -> list[Statement]:
             return []
         return BlockStatement.build_statements(self.body_node, self)
 
+    @cached_property
+    def first_statement(self) -> Statement | None:
+        if len(self.statements) == 0:
+            return None
+        return self.statements[0]
+
     def __str__(self) -> str:
         return self.signature
 
@@ -190,7 +196,6 @@ def name_node(self) -> Node:
         return node
 
     @property
-    @abstractmethod
     def name(self) -> str:
         """
         The name of the function.
@@ -200,7 +205,14 @@ def name(self) -> str:
         return name_node.text.decode()
 
     @cached_property
-    @abstractmethod
+    def exits(self) -> list[Statement]:
+        """
+        The exit statements of the function, such as return statements.
+        """
+        exits = self.statements_by_types(self.language.EXIT_STATEMENTS, recursive=True)
+        return exits
+
+    @cached_property
     def accessible_functions(self) -> list[Function]:
         funcs = []
         for file in self.file.imports:
@@ -279,7 +291,6 @@ def callees(self) -> dict[Function | FunctionDeclaration, list[Statement]]:
         return callees
 
     @cached_property
-    @abstractmethod
     def callers(self) -> dict[Function, list[Statement]]:
         """
         The functions that call this function and their corresponding call sites.
@@ -467,12 +478,12 @@ def export_cfg_dot(
         )
         graph.add_node("edge", fontname="SF Pro Rounded, system-ui", arrowhead="vee")
         graph.add_node(self.signature, label=self.dot_text, color="red")
-        if len(self.statements) == 0:
+        if self.first_statement is None:
             graph.add_node(
                 self.signature, label="No statements found", color="red", shape="box"
             )
         else:
-            graph.add_edge(self.signature, self.statements[0].signature, label="CFG")
+            graph.add_edge(self.signature, self.first_statement.signature, label="CFG")
         self.__build_cfg_graph(graph, self.statements)
 
         if with_cdg:
diff --git a/scubatrace/go/function.py b/scubatrace/go/function.py
@@ -1,7 +1,18 @@
 from __future__ import annotations
 
+from functools import cached_property
+
 from ..function import Function
-from .statement import GoBlockStatement
+from .statement import GoBlockStatement, Statement
+
 
+class GoFunction(Function, GoBlockStatement):
+    @cached_property
+    def first_statement(self) -> Statement | None:
+        for statement in self.statements:
+            if statement.node_type != "defer_statement":
+                return statement
 
-class GoFunction(Function, GoBlockStatement): ...
+    @cached_property
+    def defer_statements(self) -> list[Statement]:
+        return self.statements_by_type("defer_statement")
diff --git a/scubatrace/go/statement.py b/scubatrace/go/statement.py
@@ -1,9 +1,135 @@
 from __future__ import annotations
 
-from ..statement import BlockStatement, SimpleStatement
+from functools import cached_property
 
+from ..statement import BlockStatement, SimpleStatement, Statement
 
-class GoSimpleStatement(SimpleStatement): ...
 
+class GoSimpleStatement(SimpleStatement):
+    @cached_property
+    def post_controls(self) -> list[Statement]:
+        if self.node_type in self.language.EXIT_STATEMENTS:
+            return []
+        exits_statements = self.function.exits if self.function is not None else []
 
-class GoBlockStatement(BlockStatement): ...
+        last_defer_statement = []
+        from .function import GoFunction
+
+        if self.function is not None and isinstance(self.function, GoFunction):
+            last_defer_statement = (
+                self.function.defer_statements[-1:]
+                if len(self.function.defer_statements) > 0
+                else []
+            )
+
+        if self.node_type == "defer_statement":
+            assert isinstance(self.function, GoFunction)
+            defter_index = self.function.defer_statements.index(self)
+            return (
+                [self.function.defer_statements[defter_index - 1]]
+                if defter_index > 0
+                else exits_statements
+            )
+
+        if self.node_type in self.language.CONTINUE_STATEMENTS:
+            loop_stat = self.ancestor_by_types(self.language.LOOP_STATEMENTS)
+            return [loop_stat] if loop_stat else []
+        if self.node_type in self.language.BREAK_STATEMENTS:
+            loop_stat = self.ancestor_by_types(
+                self.language.LOOP_STATEMENTS + self.language.SWITCH_STATEMENTS
+            )
+            preorder_successor = loop_stat.preorder_successor if loop_stat else None
+            return [preorder_successor] if preorder_successor else []
+        if self.node_type in self.language.GOTO_STATEMENTS:
+            function = self.function
+            if function is not None:
+                label_name_node = self.node.child_by_field_name("label")
+                assert label_name_node is not None and label_name_node.text is not None
+                label_name = label_name_node.text.decode()
+                label_stat = function.query_oneshot(
+                    self.language.query_goto_label(label_name)
+                )
+                return [label_stat] if label_stat else []
+
+        if self.parent.node_type in self.language.LOOP_STATEMENTS:
+            # while () {last_statement;}
+            loop_stat = self.ancestor_by_types(self.language.LOOP_STATEMENTS)
+            is_last_statement = self.next_sibling is None
+            if is_last_statement:
+                return [loop_stat] if loop_stat else []
+        if self.parent.node_type in self.language.IF_STATEMENTS:
+            # if () {last_statement;} else { ...}
+            consequences = self.parent.statements_by_field_name("consequence")
+            if self in consequences:
+                is_last_consequences = consequences.index(self) == len(consequences) - 1
+                if is_last_consequences:
+                    return (
+                        [self.right_uncle_ancestor] if self.right_uncle_ancestor else []
+                    )
+
+        preorder_successor = self.preorder_successor
+        while (
+            preorder_successor is not None
+            and preorder_successor.node_type == "defer_statement"
+        ):
+            preorder_successor = preorder_successor.preorder_successor
+        if preorder_successor is None:
+            return last_defer_statement
+        if preorder_successor.node_type in self.language.EXIT_STATEMENTS:
+            return last_defer_statement
+        else:
+            return [preorder_successor]
+
+
+class GoBlockStatement(BlockStatement):
+    @cached_property
+    def post_controls(self) -> list[Statement]:
+        exits_statements = self.function.exits if self.function is not None else []
+        last_defer_statement = []
+        from .function import GoFunction
+
+        if self.function is not None and isinstance(self.function, GoFunction):
+            last_defer_statement = (
+                self.function.defer_statements[-1:]
+                if len(self.function.defer_statements) > 0
+                else []
+            )
+
+        if self.node_type in self.language.IF_STATEMENTS:
+            consequences = self.statements_by_field_name("consequence")
+            alternatives = self.statements_by_field_name("alternative")
+            nexts = []
+            if len(consequences) > 0:
+                nexts.append(consequences[0])
+            if len(alternatives) > 0:
+                nexts.append(alternatives[0])
+            elif (
+                self.preorder_successor is not None
+                and self.preorder_successor not in exits_statements
+            ):
+                nexts.append(self.preorder_successor)
+            else:
+                nexts.extend(last_defer_statement)
+            return nexts
+        if self.node_type in self.language.SWITCH_STATEMENTS:
+            if len(self.statements) > 0:
+                return [self.statements[0]]
+        if self.parent.node_type in self.language.SWITCH_STATEMENTS:
+            if self.text.strip().startswith("default:") and len(self.statements) > 0:
+                return [self.statements[0]]
+        if self.parent.node_type in self.language.LOOP_STATEMENTS:
+            # while () {last_statement;}
+            loop_stat = self.ancestor_by_types(self.language.LOOP_STATEMENTS)
+            is_last_statement = self.next_sibling is None
+            if is_last_statement:
+                return [loop_stat] if loop_stat else []
+
+        nexts = [self.statements[0]] if len(self.statements) > 0 else []
+        if self.preorder_successor is not None:
+            nexts.append(self.preorder_successor)
+
+        if len(last_defer_statement) > 0:
+            nexts = [stat for stat in nexts if stat not in exits_statements]
+            if len(nexts) == 0:
+                return last_defer_statement
+        return nexts
diff --git a/tests/samples/go/main.go b/tests/samples/go/main.go
@@ -0,0 +1,32 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+)
+
+func main() {
+	defer fmt.Println("world")
+	defer fmt.Println("hello")
+	flag.Parse()
+	defer fmt.Println("Flags parsed successfully")
+	defer func() {
+		if r := recover(); r != nil {
+			fmt.Println("Recovered from panic:", r)
+		}
+	}()
+
+	var (
+		testFlag1 = flag.String("testFlag1", "default1", "Description for testFlag1")
+		testFlag2 = flag.String("testFlag2", "default2", "Description for testFlag2")
+	)
+	if *testFlag1 != "default1" || *testFlag2 != "default2" {
+		panic("Flag values do not match expected defaults")
+	}
+	if *testFlag1 == "default1" && *testFlag2 == "default2" {
+		println("Flag values match expected defaults")
+	} else {
+		panic("Flag values do not match expected defaults")
+	}
+	println("All tests passed successfully")
+}
