regalloc: model exception-handler edges in liveness

A try body lowers to blocks whose only explicit CFG successors are the
normal-flow ones (e.g. RAISE_VARARGS -> exit); the edge from a faulting op
in the try body to the handler is implicit (established by SETUP_EXC_HANDLE
in the *predecessor*). mlir::Liveness therefore never sees body-op ->
handler, so a value that is live at the handler — or live after it — was
considered dead inside the try body, and the register allocator happily
reused its register there.

When an exception actually unwound, that register had been overwritten.
The clearest case: a `for` loop whose body contains `try/except`. The
FOR_ITER iterator is read again after the handler (next iteration), but its
register was reused for a value inside the try body, so the resumed loop
read a clobbered iterator (abort in ForIter / "object is not an iterator").
The same bug corrupted exception args across sequential try/except blocks.

Fix: in LiveAnalysis, for every block in a try body (those dominated by the
SETUP_EXC_HANDLE / SETUP_WITH try-entry successor), add the handler's
live-in set to the block's live-out before computing per-op liveness. Any
op in the try body can transfer to the handler, so values live at the
handler are live across the whole try body. Dominance keeps it bounded and
handles nesting; the pass is skipped entirely when a function has no
handlers.

This makes the iterator (and any value live past the handler) interfere
with try-body temporaries, so they no longer share a register.

Author: gf712

integration/tests/exception_types.py

@@ -2,10 +2,6 @@
 # crashing — Exception subclasses missing their own __new__ used to inherit
 # Exception::__new__ (which asserts the exact Exception type), and
 # ModuleNotFoundError dereferenced a null kwargs.
-#
-# NB: the per-type check lives in a helper called from the loop body rather than
-# inline, to avoid the (separate) FOR_ITER iterator-register clobber bug that a
-# heavy loop body triggers.
 
 builtin_exceptions = [
     BaseException, Exception, ValueError, KeyError, IndexError, TypeError,
@@ -15,7 +11,7 @@
 ]
 
 
-def check(exc_type):
+for exc_type in builtin_exceptions:
     try:
         raise exc_type("msg")
     except BaseException as e:
@@ -24,10 +20,6 @@ def check(exc_type):
         assert e.args == ("msg",), (exc_type, e.args)
 
 
-for exc_type in builtin_exceptions:
-    check(exc_type)
-
-
 # subclass relationships still hold
 try:
     raise RuntimeError("r")

integration/tests/regalloc_exception_liveness.py

@@ -0,0 +1,72 @@
+# Regression: exception-handler edges must be modelled in liveness.
+#
+# An operation inside a try body can transfer to the handler, but that edge is
+# not in the explicit CFG. When liveness ignored it, a value live across the try
+# body via the handler path (e.g. a FOR_ITER iterator, or a value used after the
+# handler) had its register reused inside the try body and was clobbered when an
+# exception actually unwound.
+
+
+# A for-loop whose body raises and catches: the iterator must survive the try
+# body. Previously clobbered (abort in FOR_ITER / "object is not an iterator").
+seen = []
+for x in [1, 2, 3]:
+    try:
+        raise ValueError("m")
+    except ValueError:
+        pass
+    seen.append(x)
+assert seen == [1, 2, 3], seen
+
+# Same over range() and over a list of types, with the exception bound.
+total = 0
+for x in range(4):
+    try:
+        raise ValueError("m")
+    except ValueError as e:
+        assert str(e) == "m"
+    total += x
+assert total == 6, total
+
+for exc in [ValueError, KeyError, RuntimeError, TypeError, NameError]:
+    try:
+        raise exc("msg")
+    except BaseException as e:
+        assert isinstance(e, exc), exc
+        assert e.args == ("msg",), (exc, e.args)
+
+# Sequential try/except in one frame must not leak the prior exception's args.
+try:
+    raise ValueError("hello")
+except ValueError as e:
+    assert e.args == ("hello",), e.args
+try:
+    raise ValueError("a", "b")
+except ValueError as e:
+    assert e.args == ("a", "b"), e.args
+
+# A recursive call whose result must survive a following try/except (the
+# original minimal miscompile repro).
+def fib(n):
+    return n if n < 2 else fib(n - 1) + fib(n - 2)
+
+
+assert fib(10) == 55
+try:
+    raise ValueError("e")
+except ValueError as e:
+    assert str(e) == "e"
+
+# Nested try/except inside a loop.
+acc = 0
+for x in [1, 2, 3]:
+    try:
+        try:
+            raise ValueError(x)
+        except KeyError:
+            pass
+    except ValueError as e:
+        acc += e.args[0]
+assert acc == 6, acc
+
+print("REGALLOC_EXCEPTION_LIVENESS_OK")

integration/tests/regalloc_exception_liveness_shapes.py

@@ -0,0 +1,158 @@
+# Regression: the exception-handler-edge liveness fix (a value live across a try
+# body via the handler path must keep its register) must hold across try/except,
+# try/finally, with, nested try, and except-cascade shapes — not just the simple
+# FOR_ITER + try/except case. Each shape loops (register pressure) and keeps a
+# value live across a multi-block / faulting try body.
+
+
+# if/else inside the try body => multi-block body; loop var + accumulator survive
+def if_else_body(flag):
+    acc = 0
+    for x in [1, 2, 3]:
+        try:
+            if flag:
+                raise ValueError("a")
+            else:
+                raise KeyError("b")
+        except ValueError:
+            pass
+        except KeyError:
+            pass
+        acc += x
+    return acc
+
+
+assert if_else_body(True) == 6, if_else_body(True)
+assert if_else_body(False) == 6, if_else_body(False)
+
+
+# a loop inside the try body; an outer value survives the inner loop + raise
+def loop_in_try():
+    out = []
+    for x in [1, 2]:
+        try:
+            for i in range(3):
+                pass
+            raise ValueError(x)
+        except ValueError:
+            pass
+        out.append(x)
+    return out
+
+
+assert loop_in_try() == [1, 2], loop_in_try()
+
+
+# try/except/finally: finally runs on both paths; loop var survives
+def try_except_finally():
+    log = []
+    for x in [1, 2]:
+        try:
+            raise ValueError(x)
+        except ValueError:
+            log.append(x)
+        finally:
+            log.append(-x)
+    return log
+
+
+assert try_except_finally() == [1, -1, 2, -2], try_except_finally()
+
+
+# try/finally with the exception path actually taken (finally on unwind)
+def try_finally_raise():
+    out = []
+    for x in [1, 2]:
+        try:
+            try:
+                raise ValueError(x)
+            finally:
+                out.append(-x)
+        except ValueError:
+            out.append(x)
+    return out
+
+
+assert try_finally_raise() == [-1, 1, -2, 2], try_finally_raise()
+
+
+class CM:
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *a):
+        return False
+
+
+# with-statement: the body may raise; loop var survives the cleanup region
+def with_body():
+    res = []
+    for x in [1, 2]:
+        try:
+            with CM():
+                raise ValueError(x)
+        except ValueError:
+            res.append(x)
+    return res
+
+
+assert with_body() == [1, 2], with_body()
+
+
+# nested try: inner clause does NOT match, exception propagates to outer
+def nested_propagate():
+    res = []
+    for x in [1, 2]:
+        try:
+            try:
+                raise ValueError(x)
+            except KeyError:
+                res.append("k")
+        except ValueError:
+            res.append(x)
+    return res
+
+
+assert nested_propagate() == [1, 2], nested_propagate()
+
+
+# multiple except clauses (type cascade); value survives the whole try
+def except_cascade(which):
+    for x in [7]:
+        try:
+            if which == 0:
+                raise ValueError(x)
+            elif which == 1:
+                raise KeyError(x)
+            else:
+                raise TypeError(x)
+        except ValueError:
+            return ("v", x)
+        except KeyError:
+            return ("k", x)
+        except TypeError:
+            return ("t", x)
+
+
+assert except_cascade(0) == ("v", 7), except_cascade(0)
+assert except_cascade(1) == ("k", 7), except_cascade(1)
+assert except_cascade(2) == ("t", 7), except_cascade(2)
+
+
+# a value defined before the try and used after the handler
+def value_after_handler():
+    out = []
+    for x in [1, 2, 3]:
+        keep = x * 10
+        try:
+            raise ValueError(x)
+        except ValueError as e:
+            got = e.args[0]
+        out.append(keep + got)
+    return out
+
+
+assert value_after_handler() == [11, 22, 33], value_after_handler()
+
+
+print("REGALLOC_EXCEPTION_LIVENESS_SHAPES_OK")

integration/tests/repr_quoting.py

@@ -19,31 +19,21 @@
 
 
 # exception repr quotes its args; str() stays the bare message
-def check_single_arg():
-    try:
-        raise ValueError("hello")
-    except ValueError as e:
-        assert repr(e) == "ValueError('hello')", repr(e)
-        assert str(e) == "hello", str(e)
-        assert e.args == ("hello",), e.args
-
-
-def check_multi_arg():
-    try:
-        raise ValueError("a", "b")
-    except ValueError as e:
-        assert repr(e) == "ValueError('a', 'b')", repr(e)
-
-
-def check_no_arg():
-    try:
-        raise ValueError
-    except ValueError as e:
-        assert repr(e) == "ValueError()", repr(e)
-
-
-check_single_arg()
-check_multi_arg()
-check_no_arg()
+try:
+    raise ValueError("hello")
+except ValueError as e:
+    assert repr(e) == "ValueError('hello')", repr(e)
+    assert str(e) == "hello", str(e)
+    assert e.args == ("hello",), e.args
+
+try:
+    raise ValueError("a", "b")
+except ValueError as e:
+    assert repr(e) == "ValueError('a', 'b')", repr(e)
+
+try:
+    raise ValueError
+except ValueError as e:
+    assert repr(e) == "ValueError()", repr(e)
 
 print("REPR_QUOTING_OK")