# -*- coding: utf-8 -*-

"""
Rewrite exceptions that are thrown and caught locally to jumps.
"""

from numba2.compiler import excmodel

from pykit.analysis import cfa
from pykit.optimizations import local_exceptions

def rewrite_local_exceptions(func, env):
    local_exceptions.run(func, env, exc_model=excmodel.ExcModel(env))

def rewrite_exceptions(func, env):
    blocks = set()
    for op in func.ops:
        if op.opcode == 'exc_throw':
            raise NotImplementedError("Exception throwing", op, func)
        if op.opcode in ('exc_catch', 'exc_setup'):
            blocks.add(op.block)
            op.delete()

    update_outdated_incoming_blocks(func, blocks)

def update_outdated_incoming_blocks(func, candidates):
    """
    Update phi nodes in blocks previously containing 'exc_catch'. 'exc_setup'
    may span many blocks, and none, or only a subset of those blocks may be
    actual predecessors.
    """
    cfg = cfa.cfg(func)
    for block in candidates:
        preds = cfg.predecessors(block)
        for op in block.leaders:
            if op.opcode == 'phi':
                blocks, values = op.args
                newblocks = [block for block in blocks if block in preds]
                newvalues = [val for block, val in zip(blocks, values)
                                     if block in preds]
                op.set_args([newblocks, newvalues])
# -*- coding: utf-8 -*-

"""
Rewrite exceptions that are thrown and caught locally to jumps.
"""

from numba2.compiler import excmodel
from pykit.optimizations import local_exceptions

def rewrite_local_exceptions(func, env):
    local_exceptions.run(func, env, exc_model=excmodel.ExcModel(env))

def rewrite_exceptions(func, env):
    for op in func.ops:
        if op.opcode == 'exc_throw':
            raise NotImplementedError("Exception throwing", op, func)
        if op.opcode in ('exc_catch', 'exc_setup'):
            op.delete()

--- 
+++ 
@@ -5,14 +5,37 @@
 """
 
 from numba2.compiler import excmodel
+
+from pykit.analysis import cfa
 from pykit.optimizations import local_exceptions
 
 def rewrite_local_exceptions(func, env):
     local_exceptions.run(func, env, exc_model=excmodel.ExcModel(env))
 
 def rewrite_exceptions(func, env):
+    blocks = set()
     for op in func.ops:
         if op.opcode == 'exc_throw':
             raise NotImplementedError("Exception throwing", op, func)
         if op.opcode in ('exc_catch', 'exc_setup'):
+            blocks.add(op.block)
             op.delete()
+
+    update_outdated_incoming_blocks(func, blocks)
+
+def update_outdated_incoming_blocks(func, candidates):
+    """
+    Update phi nodes in blocks previously containing 'exc_catch'. 'exc_setup'
+    may span many blocks, and none, or only a subset of those blocks may be
+    actual predecessors.
+    """
+    cfg = cfa.cfg(func)
+    for block in candidates:
+        preds = cfg.predecessors(block)
+        for op in block.leaders:
+            if op.opcode == 'phi':
+                blocks, values = op.args
+                newblocks = [block for block in blocks if block in preds]
+                newvalues = [val for block, val in zip(blocks, values)
+                                     if block in preds]
+                op.set_args([newblocks, newvalues])
Rewrite phis from outdated incoming exception blocks
