From 6542b65db39e8c97ff5f0119d56adfc07e487298 Mon Sep 17 00:00:00 2001
From: pca006132 <jl@m-labs.hk>
Date: Mon, 10 Jan 2022 15:21:39 +0800
Subject: [PATCH] compiler: fixed exception codegen issues

---
 artiq/compiler/ir.py                          |  8 +++++
 .../compiler/transforms/artiq_ir_generator.py | 30 ++++++++++++++++---
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/artiq/compiler/ir.py b/artiq/compiler/ir.py
index ae786f8a8..92bc4fb3a 100644
--- a/artiq/compiler/ir.py
+++ b/artiq/compiler/ir.py
@@ -1360,6 +1360,14 @@ class LandingPad(Terminator):
     def cleanup(self):
         return self.operands[0]
 
+    def erase(self):
+        self.remove_from_parent()
+        # we should erase all clauses as well
+        for block in set(self.operands):
+            block.uses.remove(self)
+            block.erase()
+        assert not any(self.uses)
+
     def clauses(self):
         return zip(self.operands[1:], self.types)
 
diff --git a/artiq/compiler/transforms/artiq_ir_generator.py b/artiq/compiler/transforms/artiq_ir_generator.py
index 5873664ed..069422d1b 100644
--- a/artiq/compiler/transforms/artiq_ir_generator.py
+++ b/artiq/compiler/transforms/artiq_ir_generator.py
@@ -626,6 +626,11 @@ class ARTIQIRGenerator(algorithm.Visitor):
             self.final_branch(raise_proxy, self.current_block)
             self.current_block = raise_proxy
 
+        if exn is not None:
+            # if we need to raise the exception in a final body, we have to
+            # lazy-evaluate the exception object to make sure that we generate
+            # it in the raise_proxy block
+            exn = exn()
         if exn is not None:
             assert loc is not None
             loc_file = ir.Constant(loc.source_buffer.name, builtins.TStr())
@@ -650,9 +655,9 @@ class ARTIQIRGenerator(algorithm.Visitor):
 
     def visit_Raise(self, node):
         if node.exc is not None and types.is_exn_constructor(node.exc.type):
-            self.raise_exn(self.alloc_exn(node.exc.type.instance), loc=self.current_loc)
+            self.raise_exn(lambda: self.alloc_exn(node.exc.type.instance), loc=self.current_loc)
         else:
-            self.raise_exn(self.visit(node.exc), loc=self.current_loc)
+            self.raise_exn(lambda: self.visit(node.exc), loc=self.current_loc)
 
     def visit_Try(self, node):
         dispatcher = self.add_block("try.dispatch")
@@ -671,6 +676,15 @@ class ARTIQIRGenerator(algorithm.Visitor):
                 final_targets.append(target)
                 final_paths.append(block)
 
+            final_exn_targets = []
+            final_exn_paths = []
+            # raise has to be treated differently
+            # we cannot follow indirectbr for local access validation, so we
+            # have to construct the control flow explicitly
+            def exception_final_branch(target, block):
+                final_exn_targets.append(target)
+                final_exn_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
@@ -714,7 +728,7 @@ 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
+            old_final_branch, self.final_branch = self.final_branch, exception_final_branch
 
         cleanup = self.add_block('handler.cleanup')
         landingpad = dispatcher.append(ir.LandingPad(cleanup))
@@ -745,6 +759,14 @@ class ARTIQIRGenerator(algorithm.Visitor):
             # Finalize and continue after try statement.
             self.final_branch = old_final_branch
 
+            for (i, (target, block)) in enumerate(zip(final_exn_targets, final_exn_paths)):
+                finalizer = self.add_block(f"finally{i}")
+                self.current_block = block
+                self.terminate(ir.Branch(finalizer))
+                self.current_block = finalizer
+                self.visit(node.finalbody)
+                self.terminate(ir.Branch(target))
+
             finalizer = self.add_block("finally")
             self.current_block = finalizer
 
@@ -997,7 +1019,7 @@ class ARTIQIRGenerator(algorithm.Visitor):
 
             old_final_branch, self.final_branch = self.final_branch, None
             old_unwind, self.unwind_target = self.unwind_target, None
-            self.raise_exn(exn_gen(*args[1:]), loc=loc)
+            self.raise_exn(lambda: exn_gen(*args[1:]), loc=loc)
         finally:
             self.current_function = old_func
             self.current_block = old_block