Make break statements trigger finally blocks

Similar to google#368, fixes an issue where break statements were bypassing
finally clauses. This change changes for and while loops to use the
checkpoint stack instead of plain goto statements so that break can
properly trigger finally handlers.

Author: trotterdylan

compiler/block.py

@@ -45,9 +45,8 @@ def __init__(self, name, alias=None):
 class Loop(object):
   """Represents a for or while loop within a particular block."""
 
-  def __init__(self, start_label, end_label):
-    self.start_label = start_label
-    self.end_label = end_label
+  def __init__(self, breakvar):
+    self.breakvar = breakvar
 
 
 class Block(object):
@@ -124,8 +123,8 @@ def free_temp(self, v):
     self.used_temps.remove(v)
     self.free_temps.add(v)
 
-  def push_loop(self):
-    loop = Loop(self.genlabel(), self.genlabel())
+  def push_loop(self, breakvar):
+    loop = Loop(breakvar)
     self.loop_stack.append(loop)
     return loop
 

compiler/block_test.py

@@ -43,9 +43,9 @@ class BlockTest(unittest.TestCase):
 
   def testLoop(self):
     b = _MakeModuleBlock()
-    loop = b.push_loop()
+    loop = b.push_loop(None)
     self.assertEqual(loop, b.top_loop())
-    inner_loop = b.push_loop()
+    inner_loop = b.push_loop(None)
     self.assertEqual(inner_loop, b.top_loop())
     b.pop_loop()
     self.assertEqual(loop, b.top_loop())

compiler/stmt.py

@@ -93,7 +93,9 @@ def visit_Break(self, node):
     if not self.block.loop_stack:
       raise util.ParseError(node, "'break' not in loop")
     self._write_py_context(node.lineno)
-    self.writer.write('goto Label{}'.format(self.block.top_loop().end_label))
+    self.writer.write_tmpl(textwrap.dedent("""\
+        $breakvar = true
+        continue"""), breakvar=self.block.top_loop().breakvar.name)
 
   def visit_ClassDef(self, node):
     # Since we only care about global vars, we end up throwing away the locals
@@ -166,7 +168,7 @@ def visit_Continue(self, node):
     if not self.block.loop_stack:
       raise util.ParseError(node, "'continue' not in loop")
     self._write_py_context(node.lineno)
-    self.writer.write('goto Label{}'.format(self.block.top_loop().start_label))
+    self.writer.write('continue')
 
   def visit_Delete(self, node):
     self._write_py_context(node.lineno)
@@ -192,40 +194,27 @@ def visit_Expr(self, node):
     self.visit_expr(node.value).free()
 
   def visit_For(self, node):
-    loop = self.block.push_loop()
-    orelse_label = self.block.genlabel() if node.orelse else loop.end_label
-    self._write_py_context(node.lineno)
-    with self.visit_expr(node.iter) as iter_expr, \
-        self.block.alloc_temp() as i, \
-        self.block.alloc_temp() as n:
-      self.writer.write_checked_call2(i, 'πg.Iter(πF, {})', iter_expr.expr)
-      self.writer.write_label(loop.start_label)
-      tmpl = textwrap.dedent("""\
-          if $n, πE = πg.Next(πF, $i); πE != nil {
-          \tisStop, exc := πg.IsInstance(πF, πE.ToObject(), πg.StopIterationType.ToObject())
-          \tif exc != nil {
-          \t\tπE = exc
-          \t\tcontinue
-          \t}
-          \tif !isStop {
-          \t\tcontinue
-          \t}
-          \tπE = nil
-          \tπF.RestoreExc(nil, nil)
-          \tgoto Label$orelse
-          }""")
-      self.writer.write_tmpl(tmpl, n=n.name, i=i.expr, orelse=orelse_label)
-      self._tie_target(node.target, n.expr)
-      self._visit_each(node.body)
-      self.writer.write('goto Label{}'.format(loop.start_label))
-
-    self.block.pop_loop()
-    if node.orelse:
-      self.writer.write_label(orelse_label)
-      self._visit_each(node.orelse)
-    # Avoid label "defined and not used" in case there's no break statements.
-    self.writer.write('goto Label{}'.format(loop.end_label))
-    self.writer.write_label(loop.end_label)
+    with self.block.alloc_temp() as i:
+      with self.visit_expr(node.iter) as iter_expr:
+        self.writer.write_checked_call2(i, 'πg.Iter(πF, {})', iter_expr.expr)
+      def testfunc(testvar):
+        with self.block.alloc_temp() as n:
+          self.writer.write_tmpl(textwrap.dedent("""\
+              if $n, πE = πg.Next(πF, $i); πE != nil {
+              \tisStop, exc := πg.IsInstance(πF, πE.ToObject(), πg.StopIterationType.ToObject())
+              \tif exc != nil {
+              \t\tπE = exc
+              \t} else if isStop {
+              \t\tπE = nil
+              \t\tπF.RestoreExc(nil, nil)
+              \t}
+              \t$testvar = !isStop
+              } else {
+              \t$testvar = true"""), n=n.name, i=i.expr, testvar=testvar.name)
+          with self.writer.indent_block():
+            self._tie_target(node.target, n.expr)
+          self.writer.write('}')
+      self._visit_loop(testfunc, node)
 
   def visit_FunctionDef(self, node):
     self._write_py_context(node.lineno + len(node.decorator_list))
@@ -357,6 +346,8 @@ def visit_Return(self, node):
     if node.value:
       with self.visit_expr(node.value) as value:
         self.writer.write('πR = {}'.format(value.expr))
+    else:
+      self.writer.write('πR = πg.None')
     self.writer.write('continue')
 
   def visit_Try(self, node):
@@ -439,26 +430,12 @@ def visit_Try(self, node):
               }"""), exc=exc.expr, tb=tb.expr)
 
   def visit_While(self, node):
-    loop = self.block.push_loop()
     self._write_py_context(node.lineno)
-    self.writer.write_label(loop.start_label)
-    orelse_label = self.block.genlabel() if node.orelse else loop.end_label
-    with self.visit_expr(node.test) as cond,\
-        self.block.alloc_temp('bool') as is_true:
-      self.writer.write_checked_call2(is_true, 'πg.IsTrue(πF, {})', cond.expr)
-      self.writer.write_tmpl(textwrap.dedent("""\
-          if !$is_true {
-          \tgoto Label$orelse_label
-          }"""), is_true=is_true.expr, orelse_label=orelse_label)
-      self._visit_each(node.body)
-      self.writer.write('goto Label{}'.format(loop.start_label))
-    if node.orelse:
-      self.writer.write_label(orelse_label)
-      self._visit_each(node.orelse)
-    # Avoid label "defined and not used" in case there's no break statements.
-    self.writer.write('goto Label{}'.format(loop.end_label))
-    self.writer.write_label(loop.end_label)
-    self.block.pop_loop()
+    def testfunc(testvar):
+      with self.visit_expr(node.test) as cond:
+        self.writer.write_checked_call2(
+            testvar, 'πg.IsTrue(πF, {})', cond.expr)
+    self._visit_loop(testfunc, node)
 
   def visit_With(self, node):
     assert len(node.items) == 1, 'multiple items in a with not yet supported'
@@ -739,6 +716,44 @@ def _visit_each(self, nodes):
     for node in nodes:
       self.visit(node)
 
+  def _visit_loop(self, testfunc, node):
+    start_label = self.block.genlabel(is_checkpoint=True)
+    else_label = self.block.genlabel(is_checkpoint=True)
+    end_label = self.block.genlabel()
+    with self.block.alloc_temp('bool') as breakvar:
+      self.block.push_loop(breakvar)
+      self.writer.write('πF.PushCheckpoint({})'.format(else_label))
+      self.writer.write('{} = false'.format(breakvar.name))
+      self.writer.write_label(start_label)
+      self.writer.write_tmpl(textwrap.dedent("""\
+          if πE != nil || πR != nil {
+          \tcontinue
+          }
+          if $breakvar {
+          \tπF.PopCheckpoint()
+          \tgoto Label$end_label
+          }"""), breakvar=breakvar.expr, end_label=end_label)
+      with self.block.alloc_temp('bool') as testvar:
+        testfunc(testvar)
+        self.writer.write_tmpl(textwrap.dedent("""\
+            if πE != nil || !$testvar {
+            \tcontinue
+            }
+            πF.PushCheckpoint($start_label)\
+            """), testvar=testvar.name, start_label=start_label)
+      self._visit_each(node.body)
+      self.writer.write('continue')
+      # End the loop so that break applies to an outer loop if present.
+      self.block.pop_loop()
+      self.writer.write_label(else_label)
+      self.writer.write(textwrap.dedent("""\
+          if πE != nil || πR != nil {
+          \tcontinue
+          }"""))
+      if node.orelse:
+        self._visit_each(node.orelse)
+      self.writer.write_label(end_label)
+
   def _write_except_block(self, label, exc, except_node):
     self._write_py_context(except_node.lineno)
     self.writer.write_label(label)

testing/try_test.py

@@ -232,3 +232,19 @@ def foo():
 
 
 assert foo() == 2
+
+
+# Break statement should not bypass finally.
+x = []
+def foo():
+  while True:
+    try:
+      x.append(1)
+      break
+    finally:
+      x.append(2)
+  x.append(3)
+
+
+foo()
+assert x == [1, 2, 3]