From 5f68cc6a21d318327bf4a9145970d78c4b61558b Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Sat, 2 Jan 2016 06:51:14 +0800
Subject: [PATCH] transforms.artiq_ir_generator: handle `raise` in `except:`
 with `finally:`.

---
 .../compiler/transforms/artiq_ir_generator.py | 40 +++++++++++++------
 1 file changed, 27 insertions(+), 13 deletions(-)

diff --git a/artiq/compiler/transforms/artiq_ir_generator.py b/artiq/compiler/transforms/artiq_ir_generator.py
index 9c6ba0259..a27afadbd 100644
--- a/artiq/compiler/transforms/artiq_ir_generator.py
+++ b/artiq/compiler/transforms/artiq_ir_generator.py
@@ -66,6 +66,10 @@ class ARTIQIRGenerator(algorithm.Visitor):
         the basic block to which ``return`` will transfer control
     :ivar unwind_target: (:class:`ir.BasicBlock` or None)
         the basic block to which unwinding will transfer control
+    :ivar final_branch: (function (target: :class:`ir.BasicBlock`, block: :class:`ir.BasicBlock)
+                         or None)
+        the function that appends to ``block`` a jump through the ``finally`` statement
+        to ``target``
 
     There is, additionally, some global state that is used to translate
     the results of analyses on AST level to IR level:
@@ -102,6 +106,7 @@ class ARTIQIRGenerator(algorithm.Visitor):
         self.continue_target = None
         self.return_target = None
         self.unwind_target = None
+        self.final_branch = None
         self.function_map = dict()
         self.variable_map = dict()
         self.method_map = defaultdict(lambda: [])
@@ -551,6 +556,11 @@ class ARTIQIRGenerator(algorithm.Visitor):
         self.append(ir.Branch(self.continue_target))
 
     def raise_exn(self, exn, loc=None):
+        if self.final_branch is not None:
+            raise_proxy = self.add_block("try.raise")
+            self.final_branch(raise_proxy, self.current_block)
+            self.current_block = raise_proxy
+
         if exn is not None:
             if loc is None:
                 loc = self.current_loc
@@ -588,29 +598,32 @@ class ARTIQIRGenerator(algorithm.Visitor):
                                              vars={ "$k": ir.TBasicBlock() })
             final_state    = self.append(ir.Alloc([], final_env_type))
             final_targets  = []
+            final_paths    = []
+
+            def final_branch(target, block):
+                block.append(ir.SetLocal(final_state, "$k", target))
+                final_targets.append(target)
+                final_paths.append(block)
 
             if self.break_target is not None:
                 break_proxy = self.add_block("try.break")
                 old_break, self.break_target = self.break_target, break_proxy
-                break_proxy.append(ir.SetLocal(final_state, "$k", old_break))
-                final_targets.append(old_break)
+                final_branch(old_break, break_proxy)
+
             if self.continue_target is not None:
                 continue_proxy = self.add_block("try.continue")
                 old_continue, self.continue_target = self.continue_target, continue_proxy
-                continue_proxy.append(ir.SetLocal(final_state, "$k", old_continue))
-                final_targets.append(old_continue)
+                final_branch(old_continue, continue_proxy)
 
             return_proxy = self.add_block("try.return")
             old_return, self.return_target = self.return_target, return_proxy
             if old_return is not None:
-                return_proxy.append(ir.SetLocal(final_state, "$k", old_return))
-                final_targets.append(old_return)
+                final_branch(old_return, return_proxy)
             else:
                 return_action = self.add_block("try.doreturn")
                 value = return_action.append(ir.GetLocal(self.current_private_env, "$return"))
                 return_action.append(ir.Return(value))
-                return_proxy.append(ir.SetLocal(final_state, "$k", return_action))
-                final_targets.append(return_action)
+                final_branch(return_action, return_proxy)
 
         body = self.add_block("try.body")
         self.append(ir.Branch(body))
@@ -635,6 +648,8 @@ class ARTIQIRGenerator(algorithm.Visitor):
                 self.continue_target = old_continue
             self.return_target = old_return
 
+            old_final_branch, self.final_branch = self.final_branch, final_branch
+
         cleanup = self.add_block('handler.cleanup')
         landingpad = dispatcher.append(ir.LandingPad(cleanup))
 
@@ -659,6 +674,8 @@ class ARTIQIRGenerator(algorithm.Visitor):
             handlers.append((handler, post_handler))
 
         if any(node.finalbody):
+            self.final_branch = old_final_branch
+
             finalizer = self.add_block("finally")
             self.current_block = finalizer
 
@@ -673,11 +690,8 @@ class ARTIQIRGenerator(algorithm.Visitor):
             final_targets.append(tail)
             final_targets.append(reraise)
 
-            if self.break_target:
-                break_proxy.append(ir.Branch(finalizer))
-            if self.continue_target:
-                continue_proxy.append(ir.Branch(finalizer))
-            return_proxy.append(ir.Branch(finalizer))
+            for block in final_paths:
+                block.append(ir.Branch(finalizer))
 
             if not body.is_terminated():
                 body.append(ir.SetLocal(final_state, "$k", tail))