Repository: cython/cython
Pull request: Avoid raising StopIteration in __next__ 
Fixes https://github.com/cython/cython/issues/3447

Problem statement:
Performance of `__next__` (efficiently signaling end of iteration)
I'm trying to write a custom collection class, and am struggling to write an iterator that is competitive with the iterator of a built-in tuple.

The benchmark:
```
timeit.timeit(lambda: sum(1 for my_tuple in iter_tuples() for item in my_tuple))
```
Previously `iter_tuples()` was yielding built-in tuples of varying lengths (but usually short: 90% of tuples have length 3 or less).
Now I'm trying to use my own `cdef class` for the tuples, but I can't get this benchmark anywhere near the original performance.

If my collection class uses `yield` inside its `__iter__` method, that results in code 30% slower than the built-in tuple.
If my collection class uses a custom `cdef class` for the iterator, that results in code 50% slower than the built-in tuple.

```
@cython.final
@cython.freelist(4)
cdef class NodeTupleIterator:
    cdef NodeTuple _c
    cdef size_t _index

    def __cinit__(self, NodeTuple c not None):
        self._c = c

    def __iter__(self):
        return self

    def __next__(self):
        if self._index == self._tuple._size:
            raise StopIteration
        node = self._c._get(self._index)
        self._index += 1
        return node

    def __length_hint__(self):
        return self._c._size - self._index
```
Profiling shows that `__next__` spends 33% of its time in `__Pyx_Raise` and another 40% in `__Pyx_Add_Traceback`.
But none of that is necessary. On the C level, `tp_iternext` can indicate the end of iteration by simply returning `NULL` without setting any exception.
Is this somehow possible from Cython? Should Cython optimize a `raise StopIteration` within `__next__` methods that way?

Currently it seems like my only solution is to avoid Cython for the iterator class :(

By the way, I've verified that the benchmark with the builtin tuple actually calls `tupleiter_next`. There doesn't seem to be any special case that gives builtin tuples an inherent advantage; the problem just seems to be the overhead of Cython generator machinery (for `yield`) or exception-handling machinery (for `__next__`).

Hints from issue discussion:
With hand-written C code for the iterator class, I end up 5% faster than the builtin tuple.
Have you tried a generator function?
I've created a standalone benchmark and attached it: [cython-bench.zip](https://github.com/cython/cython/files/4361229/cython-bench.zip)

No matter which Python way of writing an iterator I try, it's significantly slower than a builtin tuple.
```
Builtin Tuples: 100 loops, best of 5: 1.06 msec per loop
NodeTuple: 100 loops, best of 5: 2.7 msec per loop
YieldingNodeTuple: 100 loops, best of 5: 1.75 msec per loop
GenExprNodeTuple: 100 loops, best of 5: 1.88 msec per loop
```
With a tiny edit to the Cython-generated code, the NodeTupleIterator class wins over the other approaches:
```
NodeTuple (hand edited __next__): 100 loops, best of 5: 1.34 msec per loop
```
I faintly recall that there was a reason why we needed the `StopIteration` exception at the end of the generator, but I don't remember which. Some Python compatibility issue, I think. Might be worth trying again, much or that code is about a decade old or so.

Feel free to play around with it and try to tune it. Some C level profiling might yield different results these days than back then. Have a look at `GeneratorBodyDefNode`. It might also be possible to special-case `raise StopIteration` inside of an `__iter__()` `DefNode`.
> I faintly recall that there was a reason why we needed the StopIteration exception at the end of the generator, but I don't remember which. Some Python compatibility issue, I think.

The iterators for builtin tuples/lists only returns `NULL` without setting an exception. This was already the case in Python 2.7.0; so all `tp_iternext` users should be able to handle this case.
But maybe Cython sometimes generates direct calls to the `__next__` method without going through the usual logic for calling `tp_iternext`?
> maybe Cython sometimes generates direct calls to the `__next__` method without going through the usual logic for calling `tp_iternext`?

If that was the case, I'd rather see that fixed.

Commit message: Fix compiling Code.py

Patch to review:
diff --git a/Cython/Compiler/Code.pxd b/Cython/Compiler/Code.pxd
index e17e0fb1ded..59779f8bcaa 100644
--- a/Cython/Compiler/Code.pxd
+++ b/Cython/Compiler/Code.pxd
@@ -54,6 +54,7 @@ cdef class FunctionState:
     cdef public object closure_temps
     cdef public bint should_declare_error_indicator
     cdef public bint uses_error_indicator
+    cdef public bint error_without_exception
 
     @cython.locals(n=size_t)
     cpdef new_label(self, name=*)
diff --git a/Cython/Compiler/Code.py b/Cython/Compiler/Code.py
index 4695b240c42..cc0048da027 100644
--- a/Cython/Compiler/Code.py
+++ b/Cython/Compiler/Code.py
@@ -691,6 +691,7 @@ def put_code(self, globalstate):
 class FunctionState(object):
     # return_label     string          function return point label
     # error_label      string          error catch point label
+    # error_without_exception  boolean Can go to the error label without an exception (e.g. __next__ can return NULL)
     # continue_label   string          loop continue point label
     # break_label      string          loop break point label
     # return_from_error_cleanup_label string
@@ -739,6 +740,8 @@ def __init__(self, owner, names_taken=set(), scope=None):
         self.should_declare_error_indicator = False
         self.uses_error_indicator = False
 
+        self.error_without_exception = False
+
     # safety checks
 
     def validate_exit(self):
@@ -2332,8 +2335,17 @@ def put_pymethoddef_wrapper(self, entry):
             if method_noargs in method_flags:
                 # Special NOARGS methods really take no arguments besides 'self', but PyCFunction expects one.
                 func_cname = Naming.method_wrapper_prefix + func_cname
-                self.putln("static PyObject *%s(PyObject *self, CYTHON_UNUSED PyObject *arg) {return %s(self);}" % (
-                    func_cname, entry.func_cname))
+                body = "%s(self)" % entry.func_cname
+                if entry.name == "__next__":
+                    # tp_iternext can return NULL without an exception
+                    body = (
+                        "PyObject *res = %s; if (!res && !PyErr_Occurred()) PyErr_SetNone(PyExc_StopIteration); return res" %
+                        body
+                    )
+                else:
+                    body = "return %s" % body
+                self.putln("static PyObject *%s(PyObject *self, CYTHON_UNUSED PyObject *arg) { %s; }" % (
+                    func_cname, body))
         return func_cname
 
     # GIL methods
diff --git a/Cython/Compiler/FlowControl.pxd b/Cython/Compiler/FlowControl.pxd
index c876ee3b183..7d2a2c9f915 100644
--- a/Cython/Compiler/FlowControl.pxd
+++ b/Cython/Compiler/FlowControl.pxd
@@ -58,6 +58,8 @@ cdef class ControlFlow:
 
      cdef public dict assmts
 
+     cdef public int in_try_block
+
      cpdef newblock(self, ControlBlock parent=*)
      cpdef nextblock(self, ControlBlock parent=*)
      cpdef bint is_tracked(self, entry)
diff --git a/Cython/Compiler/FlowControl.py b/Cython/Compiler/FlowControl.py
index 4e0160e41b8..b64a5c1884b 100644
--- a/Cython/Compiler/FlowControl.py
+++ b/Cython/Compiler/FlowControl.py
@@ -110,6 +110,7 @@ class ControlFlow(object):
        entries     set    tracked entries
        loops       list   stack for loop descriptors
        exceptions  list   stack for exception descriptors
+       in_try_block  int  track if we're in a try...except or try...finally blcok
     """
 
     def __init__(self):
@@ -122,6 +123,7 @@ def __init__(self):
         self.exit_point = ExitBlock()
         self.blocks.add(self.exit_point)
         self.block = self.entry_point
+        self.in_try_block = 0
 
     def newblock(self, parent=None):
         """Create floating block linked to `parent` if given.
@@ -1166,7 +1168,9 @@ def visit_TryExceptStatNode(self, node):
         ## XXX: children nodes
         self.flow.block.add_child(entry_point)
         self.flow.nextblock()
+        self.flow.in_try_block += 1
         self._visit(node.body)
+        self.flow.in_try_block -= 1
         self.flow.exceptions.pop()
 
         # After exception
@@ -1226,7 +1230,9 @@ def visit_TryFinallyStatNode(self, node):
         self.flow.block = body_block
         body_block.add_child(entry_point)
         self.flow.nextblock()
+        self.flow.in_try_block += 1
         self._visit(node.body)
+        self.flow.in_try_block -= 1
         self.flow.exceptions.pop()
         if self.flow.loops:
             self.flow.loops[-1].exceptions.pop()
@@ -1245,6 +1251,8 @@ def visit_RaiseStatNode(self, node):
         if self.flow.exceptions:
             self.flow.block.add_child(self.flow.exceptions[-1].entry_point)
         self.flow.block = None
+        if self.flow.in_try_block:
+            node.in_try_block = True
         return node
 
     def visit_ReraiseStatNode(self, node):
diff --git a/Cython/Compiler/Nodes.py b/Cython/Compiler/Nodes.py
index 927e4776386..0033209e4c6 100644
--- a/Cython/Compiler/Nodes.py
+++ b/Cython/Compiler/Nodes.py
@@ -2184,6 +2184,8 @@ def assure_gil(code_path, code=code):
                 # code.put_trace_exception()
 
                 assure_gil('error')
+                if code.funcstate.error_without_exception:
+                    code.put("if (PyErr_Occurred()) ")
                 code.put_add_traceback(self.entry.qualified_name)
             else:
                 warning(self.entry.pos,
@@ -6656,11 +6658,15 @@ class RaiseStatNode(StatNode):
     #  exc_value   ExprNode or None
     #  exc_tb      ExprNode or None
     #  cause       ExprNode or None
+    #
+    # set in FlowControl
+    #  in_try_block  bool
 
     child_attrs = ["exc_type", "exc_value", "exc_tb", "cause"]
     is_terminator = True
     builtin_exc_name = None
     wrap_tuple_value = False
+    in_try_block = False
 
     def analyse_expressions(self, env):
         if self.exc_type:
@@ -6689,9 +6695,19 @@ def analyse_expressions(self, env):
                     not (exc.args or (exc.arg_tuple is not None and exc.arg_tuple.args))):
                 exc = exc.function  # extract the exception type
             if exc.is_name and exc.entry.is_builtin:
+                from . import Symtab
                 self.builtin_exc_name = exc.name
                 if self.builtin_exc_name == 'MemoryError':
                     self.exc_type = None  # has a separate implementation
+                elif (self.builtin_exc_name == 'StopIteration' and
+                        isinstance(env, Symtab.LocalScope) and env.name == "__next__" and
+                        env.parent_scope and env.parent_scope.is_c_class_scope and
+                        not self.in_try_block):
+                    # tp_iternext is allowed to return NULL without raising StopIteration.
+                    # For the sake of simplicity, only allow this to happen when not in
+                    # a try block
+                    self.exc_type = None
+
         return self
 
     nogil_check = Node.gil_error
@@ -6702,6 +6718,10 @@ def generate_execution_code(self, code):
         if self.builtin_exc_name == 'MemoryError':
             code.putln('PyErr_NoMemory(); %s' % code.error_goto(self.pos))
             return
+        elif self.builtin_exc_name == 'StopIteration' and not self.exc_type:
+            code.putln('%s;' % code.error_goto(None))
+            code.funcstate.error_without_exception = True
+            return
 
         if self.exc_type:
             self.exc_type.generate_evaluation_code(code)
diff --git a/tests/run/funcexc_iter_T228.pyx b/tests/run/funcexc_iter_T228.pyx
index 4b81166f63b..089929ad2ee 100644
--- a/tests/run/funcexc_iter_T228.pyx
+++ b/tests/run/funcexc_iter_T228.pyx
@@ -65,3 +65,88 @@ def double_raise(py_iterator):
         print(sys.exc_info()[0] is ValueError or sys.exc_info()[0])
         a = list(cy_iterator())
         print(sys.exc_info()[0] is ValueError or sys.exc_info()[0])
+
+###### Tests to do with the optimization of StopIteration to "return NULL" #######
+# we're mainly checking that
+#  1. Calling __next__ manually doesn't crash (the wrapper function adds the exception)
+#  2. if you raise a value then that value gets raised
+#  3. putting the exception in various places try...finally / try...except blocks works
+
+def call_next_directly():
+    """
+    >>> call_next_directly()
+    Traceback (most recent call last):
+    ...
+    StopIteration
+    """
+    cy_iterator().__next__()
+
+cdef class cy_iter_many_options:
+    cdef what
+    def __init__(self, what):
+        self.what = what
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self.what == "StopIteration in finally no return":
+            try:
+                raise StopIteration
+            finally:
+                print "Finally..."
+        elif self.what == "StopIteration in finally return":
+            try:
+                raise StopIteration
+            finally:
+                self.what = None
+                return "in finally"  # but will stop iterating next time
+        elif self.what == "StopIteration from finally":
+            try:
+                raise ValueError
+            finally:
+                raise StopIteration
+        elif self.what == "catch StopIteration":
+            try:
+                raise StopIteration
+            except StopIteration:
+                self.what = None
+                return "in except"  # but will stop next time
+        elif self.what == "don't catch StopIteration":
+            try:
+                raise StopIteration
+            except ValueError:
+                return 0
+        elif self.what == "StopIteration from except":
+            try:
+                raise ValueError
+            except ValueError:
+                raise StopIteration
+        elif self.what == "StopIteration with value":
+            raise StopIteration("I'm a value!")
+        elif self.what is None:
+            raise StopIteration
+        else:
+            raise ValueError("self.what didn't match anything")
+
+def test_cy_iter_many_options(option):
+    """
+    >>> test_cy_iter_many_options("StopIteration in finally no return")
+    Finally...
+    []
+    >>> test_cy_iter_many_options("StopIteration in finally return")
+    ['in finally']
+    >>> test_cy_iter_many_options("StopIteration from finally")
+    []
+    >>> test_cy_iter_many_options("catch StopIteration")
+    ['in except']
+    >>> test_cy_iter_many_options("don't catch StopIteration")
+    []
+    >>> try:
+    ...     cy_iter_many_options("StopIteration with value").__next__()
+    ... except StopIteration as e:
+    ...     print(e.args)
+    ("I'm a value!",)
+    """
+    return list(cy_iter_many_options(option))
+


Code review comments:
diff_hunk: @@ -2332,8 +2335,17 @@ def put_pymethoddef_wrapper(self, entry):
             if method_noargs in method_flags:
                 # Special NOARGS methods really take no arguments besides 'self', but PyCFunction expects one.
                 func_cname = Naming.method_wrapper_prefix + func_cname
-                self.putln("static PyObject *%s(PyObject *self, CYTHON_UNUSED PyObject *arg) {return %s(self);}" % (
-                    func_cname, entry.func_cname))
+                body = "%s(self)" % entry.func_cname
path: Cython/Compiler/Code.py
text: @user1:
Let's not call this `body`. It confused me since I was excepting a function body, until I realised it was only a function call, and then it started changing meaning along the way.

Also, there is no reason why all of this has to go into a single C code line, right? Chaining an `if` in C with unrelated code seems very brittle. Just generate the code sequentially, as usual, and all problems should go away.

diff_hunk: @@ -6689,9 +6695,19 @@ def analyse_expressions(self, env):
                     not (exc.args or (exc.arg_tuple is not None and exc.arg_tuple.args))):
                 exc = exc.function  # extract the exception type
             if exc.is_name and exc.entry.is_builtin:
+                from . import Symtab
                 self.builtin_exc_name = exc.name
                 if self.builtin_exc_name == 'MemoryError':
                     self.exc_type = None  # has a separate implementation
+                elif (self.builtin_exc_name == 'StopIteration' and
+                        isinstance(env, Symtab.LocalScope) and env.name == "__next__" and
path: Cython/Compiler/Nodes.py
text: @user1:
Is this `isinstance()` check needed? What else should there be when the parent is a C class scope?

@author:
I could be wrong because I haven't tried it but potentially:

```
cdef class C
   class Nested:
      raise StopIteration
```
?

@user1:
Can we have an `is_local_scope` attribute then? That would seem more consistent with what we do pretty much everywhere else. (And it would avoid the local import.)

diff_hunk: @@ -2184,6 +2184,8 @@ def assure_gil(code_path, code=code):
                 # code.put_trace_exception()
 
                 assure_gil('error')
+                if code.funcstate.error_without_exception:
+                    code.put("if (PyErr_Occurred()) ")
                 code.put_add_traceback(self.entry.qualified_name)
path: Cython/Compiler/Nodes.py
text: @user1:
This is making the non-local assumption that `put_add_traceback()` issues a single command. If that ever stops being true, this code may lead to subtle bugs. Curly braces are cheap in C, let's just use them.

Also, we may be able to use `__Pyx_PyErr_Occurred()` here, which avoids a lookup of the thread state (if it has been looked up before, so not sure if this always applies).

An alternative might be a C flag variable that remembers when no exception was set.

@author:
I considered C flag variables and decided it was a little more complicated, although probably more efficient. There's a third option that might work though:

currently the code generated is:
```
error_label:
Py_XDECREF(cleanup_code)
__Pyx_Add_Traceback(....)
```

If we could reorder it to:

```
error_label:
__Pyx_AddTraceback(...)
stop_iteration_label:
__Pyx_XDECREF(variable_cleanup)
```

then that'd be very simple. I wasn't sure if there was a good reason to add the traceback after the cleanup though.

@author:
I couldn't immediately see how to use `__Pyx_PyErr_Occurred` (I think I need to generate a threadstate temp?) so I've left this suggestion for the moment.

But added the curly braces

@author:
I've switched to using a C flag variable now.

On a quick microbenchmark (essentially a simplified version of `list(reimplementation_of_range(10))` I can't see any difference between a flag variable and `PyErr_Occurred()`. Both of those give a 30% speed up vs the current Cython master and a 90% speedup vs plain Python.

Therefore the change is worthwhile I think, but it isn't worth worrying about exactly how we detect the lack of error.

diff_hunk: @@ -58,6 +58,8 @@ cdef class ControlFlow:
 
      cdef public dict assmts
 
+     cdef public int in_try_block
path: Cython/Compiler/FlowControl.pxd
text: @user1:
This may seem excessive, but I don't see a reason to apply arbitrary limits.
```suggestion
     cdef public Py_ssize_t in_try_block
```