There can be only 1: don't use a list for prev/next

DinoV · DinoV · commit 5b22166e1f5c · 2021-01-28T11:14:38.000-08:00
Summary: This is just a minor cleanup of something that has bugged me about the compiler for a while - for some reason they use a list of blocks for the next/prev block.  But they expect the list to only have at maximum 1 item.  So now we just store the blocks directly and init them to None.

Test Plan: ./python -m test.test_compiler
diff --git a/compiler/pyassem.py b/compiler/pyassem.py
@@ -194,8 +194,8 @@ def __init__(self, label=''):
         self.outEdges = set()
         self.label = label
         self.bid = None
-        self.next = []
-        self.prev = []
+        self.next = None
+        self.prev = None
         self.returns = False
         self.offset = 0
         self.seen = False # visited during stack depth calculation
@@ -205,7 +205,7 @@ def __repr__(self):
         data = []
         data.append(f'id={self.bid}')
         if self.next:
-            data.append(f'next={self.next[0].bid}')
+            data.append(f'next={self.next.bid}')
         extras = ", ".join(data)
         if self.label:
             return f"<block {self.label} {extras}>"
@@ -230,10 +230,10 @@ def addOutEdge(self, block):
         self.outEdges.add(block)
 
     def addNext(self, block):
-        self.next.append(block)
-        assert len(self.next) == 1, map(str, self.next)
-        block.prev.append(self)
-        assert len(block.prev) == 1, map(str, block.prev)
+        assert self.next is None, next
+        self.next = block
+        assert block.prev is None, block.prev
+        block.prev = self
 
     _uncond_transfer = ('RETURN_VALUE', 'RAISE_VARARGS',
                         'JUMP_ABSOLUTE', 'JUMP_FORWARD', 'CONTINUE_LOOP',
@@ -250,11 +250,11 @@ def has_return(self):
         return self.insts and self.insts[-1].opname == "RETURN_VALUE"
 
     def get_children(self):
-        return list(self.outEdges) + self.next
+        return list(self.outEdges) + [self.next]
 
     def get_followers(self):
         """Get the whole list of followers, including the next block."""
-        followers = set(self.next)
+        followers = {self.next}
         # Blocks that must be emitted *after* this one, because of
         # bytecode offsets (e.g. relative jumps) pointing to them.
         for inst in self.insts:
@@ -372,6 +372,8 @@ def dump(self, io=None):
             sys.stdout = save
 
     def stackdepth_walk(self, block, depth, maxdepth):
+        assert block is not None
+
         if block.seen or block.startdepth >= depth:
             return maxdepth
         block.seen = True
@@ -407,8 +409,7 @@ def stackdepth_walk(self, block, depth, maxdepth):
                     return maxdepth
 
         if block.next:
-            assert len(block.next) == 1
-            maxdepth = self.stackdepth_walk(block.next[0], depth, maxdepth)
+            maxdepth = self.stackdepth_walk(block.next, depth, maxdepth)
 
         block.seen = False
         return maxdepth
diff --git a/test_compiler/test_graph.py b/test_compiler/test_graph.py
@@ -19,7 +19,7 @@ class GraphTests(CompilerTest):
 
     def format_graph(self, graph):
         if graph.next:
-            return f"Block({repr(graph.label)}, {self.format_graph(graph.next[0])})"
+            return f"Block({repr(graph.label)}, {self.format_graph(graph.next)})"
         return f"Block({repr(graph.label)})"
 
     def assert_graph_equal(self, graph, expected):
@@ -34,10 +34,10 @@ def assert_graph_equal(self, graph, expected):
     def assert_graph_equal_worker(self, compiled, expected):
         self.assertEqual(compiled.label, expected.label)
         if expected.next:
-            self.assertEqual(len(compiled.next), 1)
-            self.assert_graph_equal_worker(compiled.next[0], expected.next)
+            self.assertIsNotNone(compiled.next)
+            self.assert_graph_equal_worker(compiled.next, expected.next)
         else:
-            self.assertEqual(compiled.next, [])
+            self.assertEqual(compiled.next, None)
 
     def get_child_graph(self, graph, name):
         for block in graph.ordered_blocks:
