Fix exceptions

Author: JukkaL

mypyc/irbuild/nonlocalcontrol.py

@@ -113,6 +113,7 @@ def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
 
         true, false = BasicBlock(), BasicBlock()
         stop_iter_reg = builder.fn_info.generator_class.stop_iter_value_reg
+        assert stop_iter_reg is not None
 
         builder.add(Branch(stop_iter_reg, true, false, Branch.IS_ERROR))
 
@@ -124,9 +125,9 @@ def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
         builder.builder.pop_error_handler()
 
         builder.activate_block(false)
-        # Store return value via caller-provided pointer instead of raising
-        # an exception. This fast path can only be used when the caller is a
-        # native function.
+        # The fast path is to store return value via caller-provided pointer
+        # instead of raising an exception. This can only be used when the
+        # caller is a native function.
         builder.add(SetMem(object_rprimitive, stop_iter_reg, value))
         builder.add(Return(Integer(0, object_rprimitive)))
 

mypyc/irbuild/statement.py

@@ -103,6 +103,7 @@
     get_exc_info_op,
     get_exc_value_op,
     keep_propagating_op,
+    propagate_if_error_op,
     raise_exception_op,
     reraise_exception_op,
     restore_exc_info_op,
@@ -965,7 +966,8 @@ def emit_yield_from_or_await(
         obj = builder.read(iter_reg)
         nn = builder.none_object()
         stop_iter_val = Register(object_rprimitive)
-        builder.assign(stop_iter_val, Integer(0, object_rprimitive), line)
+        err = builder.add(LoadErrorValue(object_rprimitive, undefines=True))
+        builder.assign(stop_iter_val, err, line)
         ptr = builder.add(LoadAddress(object_pointer_rprimitive, stop_iter_val))
         m = MethodCall(obj, helper_method, [nn, nn, nn, nn, ptr], line)
         # Generators have custom error handling, so disable normal error handling.
@@ -979,6 +981,7 @@ def emit_yield_from_or_await(
 
     builder.activate_block(stop_block)
     if fast_path:
+        builder.primitive_op(propagate_if_error_op, [stop_iter_val], line)
         builder.assign(result, stop_iter_val, line)
     else:
         # Try extracting a return value from a StopIteration and return it.

mypyc/lower/misc_ops.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-from mypyc.ir.ops import GetElementPtr, LoadMem, Value
-from mypyc.ir.rtypes import PyVarObject, c_pyssize_t_rprimitive
+from mypyc.ir.ops import ComparisonOp, GetElementPtr, Integer, LoadMem, Value
+from mypyc.ir.rtypes import PyVarObject, c_pyssize_t_rprimitive, object_rprimitive
 from mypyc.irbuild.ll_builder import LowLevelIRBuilder
 from mypyc.lower.registry import lower_primitive_op
 
@@ -10,3 +10,8 @@
 def var_object_size(builder: LowLevelIRBuilder, args: list[Value], line: int) -> Value:
     elem_address = builder.add(GetElementPtr(args[0], PyVarObject, "ob_size"))
     return builder.add(LoadMem(c_pyssize_t_rprimitive, elem_address))
+
+
+@lower_primitive_op("propagate_if_error")
+def propagate_if_error_op(builder: LowLevelIRBuilder, args: list[Value], line: int) -> Value:
+    return builder.add(ComparisonOp(args[0], Integer(0, object_rprimitive), ComparisonOp.NEQ))

mypyc/primitives/exc_ops.py

@@ -4,7 +4,7 @@
 
 from mypyc.ir.ops import ERR_ALWAYS, ERR_FALSE, ERR_NEVER
 from mypyc.ir.rtypes import bit_rprimitive, exc_rtuple, object_rprimitive, void_rtype
-from mypyc.primitives.registry import custom_op
+from mypyc.primitives.registry import custom_op, custom_primitive_op
 
 # If the argument is a class, raise an instance of the class. Otherwise, assume
 # that the argument is an exception object, and raise it.
@@ -62,6 +62,16 @@
     error_kind=ERR_FALSE,
 )
 
+# If argument is NULL, propagate currently raised exception (in this case
+# an exception must have been raised). If this can be used, it's faster
+# than using PyErr_Occurred().
+propagate_if_error_op = custom_primitive_op(
+    "propagate_if_error",
+    arg_types=[object_rprimitive],
+    return_type=bit_rprimitive,
+    error_kind=ERR_FALSE,
+)
+
 # Catches a propagating exception and makes it the "currently
 # handled exception" (by sticking it into sys.exc_info()). Returns the
 # exception that was previously being handled, which must be restored