Fixed dataflow graph.

src/dataflow.py

@@ -4,19 +4,23 @@ from statement import Block
 
 
 class BasicBlock(Block):
-    edges_to = []
-    edges_from = []
+    def __init__(self, statements=[]):
+        Block.__init__(self, statements)
+        self.edges_to = []
+        self.edges_from = []
 
-    dominates = []
-    dominated_by = []
+        self.dominates = []
+        self.dominated_by = []
 
     def add_edge_to(self, block):
-        self.edges_to.append(block)
-        block.edges_from.append(self)
+        if block not in self.edges_to:
+            self.edges_to.append(block)
+            block.edges_from.append(self)
 
     def set_dominates(self, block):
-        self.dominates.append(block)
-        block.dominated_by.append(self)
+        if block not in self.dominates:
+            self.dominates.append(block)
+            block.dominated_by.append(self)
 
     def get_gen(self):
         pass
@@ -74,7 +78,7 @@ def find_basic_blocks(statements):
 def generate_flow_graph(blocks):
     """Add flow graph edge administration of an ordered sequence of basic
     blocks."""
-    for b in blocks:
+    for i, b in enumerate(blocks):
         last_statement = b[-1]
 
         if last_statement.is_jump():
@@ -86,6 +90,12 @@ def generate_flow_graph(blocks):
                 if other[0].is_label(target):
                     b.add_edge_to(other)
 
+            # A branch instruction also creates an edge to the next block
+            if last_statement.is_branch() and i < len(blocks) - 1:
+                b.add_edge_to(blocks[i + 1])
+        elif i < len(blocks) - 1:
+            b.add_edge_to(blocks[i + 1])
+
 
 def generate_dominator_tree(nodes):
     """Add dominator administration to the given flow graph nodes."""

tests/test_dataflow.py

@@ -1,7 +1,8 @@
 import unittest
 
-from src.statement import Statement as S, Block as B
-from src.dataflow import find_leaders, find_basic_blocks
+from src.statement import Statement as S
+from src.dataflow import BasicBlock as B, find_leaders, find_basic_blocks, \
+        generate_flow_graph
 
 
 class TestDataflow(unittest.TestCase):
@@ -23,3 +24,23 @@ class TestDataflow(unittest.TestCase):
                 [B(s[:2]).statements, B(s[2:4]).statements, \
                  B(s[4:]).statements])
 
+    def test_generate_flow_graph_simple(self):
+        b1 = B([S('command', 'foo'), S('command', 'j', 'b2')])
+        b2 = B([S('label', 'b2'), S('command', 'bar')])
+        generate_flow_graph([b1, b2])
+
+        self.assertEqual(b1.edges_to, [b2])
+        self.assertEqual(b2.edges_from, [b1])
+
+    def test_generate_flow_graph_branch(self):
+        b1 = B([S('command', 'foo'), S('command', 'beq', '$1', '$2', 'b3')])
+        b2 = B([S('command', 'bar')])
+        b3 = B([S('label', 'b3'), S('command', 'baz')])
+        generate_flow_graph([b1, b2, b3])
+
+        self.assertIn(b2, b1.edges_to)
+        self.assertIn(b3, b1.edges_to)
+        self.assertEqual(b2.edges_from, [b1])
+        self.assertEqual(b2.edges_to, [b3])
+        self.assertIn(b1, b3.edges_from)
+        self.assertIn(b2, b3.edges_from)