Consolidate exit detection and add tests

Move the bit of exit detection logic that was in _AnnotateNode to
_MarkExitNodes(). Even if it's a bit less efficient, it's easier to
follow.

Also, fix comments in _MarkExitNodes() and add a unit test.

Author: lupino3

callgraph/core.py

@@ -83,11 +83,22 @@ def GetOrCreateNode(self, name):
         return node
 
     def _MarkExitNodes(self):
-        # Algorithm:
-        # visit all nodes via goto or nested connections, starting from __begin__
-        # mark the node as exit nodes if it contains 
+        # A node is an exit node if:
+        # 1. it contains an "exit" command with no target
+        # or
+        # 2. it's reached from the starting node via "goto" or "nested" connections
+        #    and it contains an exit command or a "goto eof" command.
+
+        # Identify all nodes with an exit command with no targets.
+        for node in self.nodes.values():
+            all_commands = set(itertools.chain.from_iterable(line.commands for line in node.code))
+            exit_cmd = Command("exit", "")
+            if exit_cmd in all_commands:
+                node.is_exit_node = True
+
+        # Visit the call graph to find nodes satisfying condition #2.
         q = [self.first_node]
-        visited = set()
+        visited = set()     # Used to avoid loops, since the call graph is not acyclic.
 
         while q:
             cur = q.pop()
@@ -99,7 +110,7 @@ def _MarkExitNodes(self):
             else:
                 all_commands = itertools.chain.from_iterable(line.commands for line in cur.code)
                 for command in all_commands:
-                    if command[0] == "exit" and command[1] == "" or command[0] == "goto" and command[1] == "eof":
+                    if command[0] == "exit" or (command[0] == "goto" and command[1] == "eof"):
                         cur.is_exit_node = True
                         break
 
@@ -166,11 +177,6 @@ def _AnnotateNode(self, node):
 
                 if command == "exit" and target == "":
                     line.terminating = True
-
-                    # Note that exit node detection is implemented in _MarkExitNodes, but that algorithm only
-                    # follows "nested" and "goto" nodes, therefore we need to mark nodes that contain naked
-                    # "exit" statements here, since they are terminating even when reached via "call" statements.
-                    node.is_exit_node = True
 
     @staticmethod
     def Build(input_file, log_file=sys.stderr):

tests/test_core.py

@@ -284,6 +284,23 @@ def test_call_nested_eof(self):
         self.assertTrue(eof.is_exit_node)
         self.assertTrue(eof.is_last_node)
 
+    def test_multiple_exit_nodes(self):
+        code = """
+        exit
+        :foo
+        exit
+        """.split("\n")
+        call_graph = CallGraph.Build(code, self.devnull)
+        self.assertEqual(2, len(call_graph.nodes))
+
+        begin = call_graph.nodes["__begin__"]
+        self.assertTrue(begin.is_exit_node)
+        self.assertFalse(begin.is_last_node)
+
+        foo = call_graph.nodes["foo"]
+        self.assertTrue(foo.is_exit_node)
+        self.assertTrue(foo.is_last_node)
+    
 
 if __name__ == "__main__":
     unittest.main()