Add comments and minor refactoring

python · JukkaL · Jul 8, 2025 · Jun 19, 2025 · Jun 19, 2025 · Jul 7, 2025
commit 22dbf453281c3e62f2b399022d211d2ef2a55dae
diff --git a/mypyc/irbuild/context.py b/mypyc/irbuild/context.py
@@ -167,6 +167,9 @@ def __init__(self, ir: ClassIR) -> None:
         # Holds the arg passed to send
         self.send_arg_reg: Value | None = None
 
+        # Holds the PyObject ** pointer through which return value can be passed
+        # instead of raising StopIteration(ret_value) (only if not NULL). This
+        # is used for faster native-to-native calls.
         self.stop_iter_value_reg: Value | None = None
 
         # The switch block is used to decide which instruction to go using the value held in the

diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py
@@ -376,8 +376,10 @@ def setup_env_for_generator_class(builder: IRBuilder) -> None:
     exc_tb = builder.add_local(Var("traceback"), object_rprimitive, is_arg=True)
     # TODO: Use the right type here instead of object?
     exc_arg = builder.add_local(Var("arg"), object_rprimitive, is_arg=True)
+
+    # Parameter that can used to pass a pointer which can used instead of
+    # raising StopIteration(value). If the value is NULL, this won't be used.
     stop_iter_value_arg = builder.add_local(Var("stop_iter_ptr"), object_pointer_rprimitive, is_arg=True)
-    # TODO: do something with stop_iter_value_arg
 
     cls.exc_regs = (exc_type, exc_val, exc_tb)
     cls.send_arg_reg = exc_arg

diff --git a/mypyc/irbuild/nonlocalcontrol.py b/mypyc/irbuild/nonlocalcontrol.py
@@ -113,20 +113,22 @@ def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
         # (???).
 
         true, false = BasicBlock(), BasicBlock()
-        r = builder.fn_info.generator_class.stop_iter_value_reg
+        stop_iter_reg = builder.fn_info.generator_class.stop_iter_value_reg
 
-        builder.add(Branch(r, true, false, Branch.IS_ERROR))
+        builder.add(Branch(stop_iter_reg, true, false, Branch.IS_ERROR))
 
         builder.activate_block(true)
-
+        # The default/slow path is to raise a StopIteration exception with
+        # return value.
         builder.call_c(set_stop_iteration_value, [value], NO_TRACEBACK_LINE_NO)
         builder.add(Unreachable())
         builder.builder.pop_error_handler()
 
         builder.activate_block(false)
-
-        # Assign to provided pointer instead of raising an exception
-        builder.add(SetMem(object_rprimitive, r, value))
+        # 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.
+        builder.add(SetMem(object_rprimitive, stop_iter_reg, value))
         builder.add(Return(Integer(0, object_rprimitive)))
 
 class CleanupNonlocalControl(NonlocalControl):

diff --git a/mypyc/irbuild/prepare.py b/mypyc/irbuild/prepare.py
@@ -221,6 +221,7 @@ def create_generator_class_if_needed(
                 RuntimeArg("value", object_rprimitive),
                 RuntimeArg("traceback", object_rprimitive),
                 RuntimeArg("arg", object_rprimitive),
+                # If not NULL, store value via this instead of raising StopIteration
                 RuntimeArg("stop_iter_ptr", object_pointer_rprimitive),
             ),
             object_rprimitive,

diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py
@@ -958,6 +958,9 @@ def emit_yield_from_or_await(
 
     if isinstance(iter_reg.type, RInstance) and iter_reg.type.class_ir.has_method(helper_method):
         # Second fast path optimization: call helper directly (see also comment above).
+        #
+        # Calling a generated generator, so avoid raising StopIteration by passing
+        # an extra PyObject ** argument to helper where the stop iteration value is stored.
         fast_path = True
         obj = builder.read(iter_reg)
         nn = builder.none_object()
@@ -974,12 +977,12 @@ def emit_yield_from_or_await(
 
     builder.add(Branch(_y_init, stop_block, main_block, Branch.IS_ERROR))
 
-    # Try extracting a return value from a StopIteration and return it.
-    # If it wasn't, this reraises the exception.
     builder.activate_block(stop_block)
     if fast_path:
         builder.assign(result, stop_iter_val, line)
     else:
+        # Try extracting a return value from a StopIteration and return it.
+        # If it wasn't, this reraises the exception.
         builder.assign(result, builder.call_c(check_stop_op, [], line), line)
     # Clear the spilled iterator/coroutine so that it will be freed.
     # Otherwise, the freeing of the spilled register would likely be delayed.