improve perf

include/pocketpy/interpreter/bindings.h

@@ -2,10 +2,32 @@
 
 #include "pocketpy/pocketpy.h"
 
-bool generator__next__(int argc, py_Ref argv);
-bool array2d_like_iterator__next__(int argc, py_Ref argv);
-bool list_iterator__next__(int argc, py_Ref argv);
-bool tuple_iterator__next__(int argc, py_Ref argv);
-bool dict_items__next__(int argc, py_Ref argv);
-bool range_iterator__next__(int argc, py_Ref argv);
-bool str_iterator__next__(int argc, py_Ref argv);
+/* Exception-free `__next__` of the builtin iterators. `py_next()` calls these
+ * directly so that exhausting a builtin iterator does not construct a
+ * `StopIteration` object.
+ *   1: a value was produced into `py_retval()`
+ *   0: the iterator is exhausted; `py_retval()` holds the `StopIteration`
+ *      value, or `nil` if there is none
+ *  -1: an error occurred and an exception was set */
+int generator__iternext(py_Ref self);
+int array2d_like_iterator__iternext(py_Ref self);
+int list_iterator__iternext(py_Ref self);
+int tuple_iterator__iternext(py_Ref self);
+int dict_items__iternext(py_Ref self);
+int range_iterator__iternext(py_Ref self);
+int str_iterator__iternext(py_Ref self);
+
+/// Raise `StopIteration` for an `__iternext` result of `0`.
+/// `py_retval()` must hold the value to carry, or `nil` for no value.
+bool pk__raise_stopiteration() PY_RAISE;
+
+/// Define the `__next__` magic method as the exception-based wrapper of
+/// `name##__iternext`. Only the owning type binds it, so it stays file-local.
+#define PK_DEFINE_NEXT_WRAPPER(name)                                                               \
+    static bool name##__next__(int argc, py_Ref argv) {                                            \
+        PY_CHECK_ARGC(1);                                                                          \
+        int res = name##__iternext(argv);                                                          \
+        if(res == -1) return false;                                                                \
+        if(res == 0) return pk__raise_stopiteration();                                             \
+        return true;                                                                               \
+    }

include/pocketpy/pocketpy.h

@@ -587,7 +587,10 @@ PK_API bool py_hash(py_Ref, py_i64* out) PY_RAISE;
 /// Get the iterator of the object.
 PK_API bool py_iter(py_Ref) PY_RAISE PY_RETURN;
 /// Get the next element from the iterator.
-/// 1: success, 0: StopIteration, -1: error
+/// 1: a value was produced into `py_retval()`
+/// 0: the iterator is exhausted; `py_retval()` holds the `StopIteration` value,
+///    or `nil` if there is none
+/// -1: error
 PK_API int py_next(py_Ref) PY_RAISE PY_RETURN;
 /// Python equivalent to `str(val)`.
 PK_API bool py_str(py_Ref val) PY_RAISE PY_RETURN;

src/bindings/py_array.c

@@ -1,6 +1,7 @@
 #include "pocketpy/pocketpy.h"
 #include "pocketpy/objects/object.h"
 #include "pocketpy/objects/iterator.h"
+#include "pocketpy/interpreter/bindings.h"
 #include "pocketpy/interpreter/vm.h"
 
 int pk_arrayview(py_Ref self, py_TValue** p) {
@@ -57,27 +58,31 @@ bool pk_arraycontains(py_Ref self, py_Ref val) {
     return true;
 }
 
-bool list_iterator__next__(int argc, py_Ref argv) {
-    PY_CHECK_ARGC(1);
-    list_iterator* ud = py_touserdata(argv);
+PK_DEFINE_NEXT_WRAPPER(list_iterator)
+
+int list_iterator__iternext(py_Ref self) {
+    list_iterator* ud = py_touserdata(self);
     if(ud->index < ud->vec->length) {
         py_TValue* res = c11__at(py_TValue, ud->vec, ud->index);
         py_assign(py_retval(), res);
         ud->index++;
-        return true;
+        return 1;
     }
-    return StopIteration();
+    py_newnil(py_retval());
+    return 0;
 }
 
-bool tuple_iterator__next__(int argc, py_Ref argv) {
-    PY_CHECK_ARGC(1);
-    tuple_iterator* ud = py_touserdata(argv);
+PK_DEFINE_NEXT_WRAPPER(tuple_iterator)
+
+int tuple_iterator__iternext(py_Ref self) {
+    tuple_iterator* ud = py_touserdata(self);
     if(ud->index < ud->length) {
         py_assign(py_retval(), ud->p + ud->index);
         ud->index++;
-        return true;
+        return 1;
     }
-    return StopIteration();
+    py_newnil(py_retval());
+    return 0;
 }
 
 py_Type pk_list_iterator__register() {

src/bindings/py_range.c

@@ -3,6 +3,7 @@
 #include "pocketpy/common/utils.h"
 #include "pocketpy/objects/object.h"
 #include "pocketpy/interpreter/vm.h"
+#include "pocketpy/interpreter/bindings.h"
 
 typedef struct Range {
     py_i64 start;
@@ -68,17 +69,19 @@ static bool range_iterator__new__(int argc, py_Ref argv) {
     return true;
 }
 
-bool range_iterator__next__(int argc, py_Ref argv) {
-    PY_CHECK_ARGC(1);
-    RangeIterator* ud = py_touserdata(argv);
-    if(ud->range.step > 0) {
-        if(ud->current >= ud->range.stop) return StopIteration();
-    } else {
-        if(ud->current <= ud->range.stop) return StopIteration();
+PK_DEFINE_NEXT_WRAPPER(range_iterator)
+
+int range_iterator__iternext(py_Ref self) {
+    RangeIterator* ud = py_touserdata(self);
+    bool exhausted = ud->range.step > 0 ? ud->current >= ud->range.stop
+                                        : ud->current <= ud->range.stop;
+    if(exhausted) {
+        py_newnil(py_retval());
+        return 0;
     }
     py_newint(py_retval(), ud->current);
     ud->current += ud->range.step;
-    return true;
+    return 1;
 }
 
 py_Type pk_range_iterator__register() {

src/bindings/py_str.c

@@ -5,6 +5,7 @@
 #include "pocketpy/objects/object.h"
 #include "pocketpy/interpreter/vm.h"
 #include "pocketpy/common/sstream.h"
+#include "pocketpy/interpreter/bindings.h"
 #include <stdbool.h>
 
 c11_string* pk_tostr(py_Ref self) {
@@ -680,17 +681,21 @@ py_Type pk_str__register() {
     return type;
 }
 
-bool str_iterator__next__(int argc, py_Ref argv) {
-    PY_CHECK_ARGC(1);
-    int* ud = py_touserdata(&argv[0]);
+PK_DEFINE_NEXT_WRAPPER(str_iterator)
+
+int str_iterator__iternext(py_Ref self) {
+    int* ud = py_touserdata(self);
     int size;
-    const char* data = py_tostrn(py_getslot(argv, 0), &size);
-    if(*ud == size) return StopIteration();
+    const char* data = py_tostrn(py_getslot(self, 0), &size);
+    if(*ud == size) {
+        py_newnil(py_retval());
+        return 0;
+    }
     int start = *ud;
     int len = c11__u8_header(data[*ud], false);
     *ud += len;
     py_newstrv(py_retval(), (c11_sv){data + start, len});
-    return true;
+    return 1;
 }
 
 py_Type pk_str_iterator__register() {

src/interpreter/ceval.c

@@ -859,9 +859,8 @@ __NEXT_STEP:
             if(res) {
                 return RES_YIELD;
             } else {
-                assert(self->last_retval.type == tp_StopIteration);
-                BaseException* ud = py_touserdata(py_retval());
-                py_ObjectRef value = &ud->args;
+                // `py_next` leaves the StopIteration value in `py_retval()`
+                py_Ref value = py_retval();
                 if(py_isnil(value)) value = py_None();
                 *TOP() = *value;  // [iter] -> [retval]
                 DISPATCH_JUMP((int16_t)byte.arg);
@@ -929,7 +928,6 @@ __NEXT_STEP:
                 PUSH(py_retval());
                 DISPATCH();
             } else {
-                assert(self->last_retval.type == tp_StopIteration);
                 POP();  // [iter] -> []
                 DISPATCH_JUMP((int16_t)byte.arg);
             }

src/interpreter/generator.c

@@ -1,7 +1,9 @@
 #include "pocketpy/interpreter/generator.h"
 #include "pocketpy/interpreter/frame.h"
 #include "pocketpy/interpreter/vm.h"
+#include "pocketpy/interpreter/bindings.h"
 #include "pocketpy/objects/base.h"
+#include "pocketpy/objects/exception.h"
 #include "pocketpy/pocketpy.h"
 #include <stdbool.h>
 #include <assert.h>
@@ -21,12 +23,16 @@ void Generator__dtor(Generator* ud) {
     if(ud->frame) Frame__delete(ud->frame);
 }
 
-bool generator__next__(int argc, py_Ref argv) {
-    PY_CHECK_ARGC(1);
-    Generator* ud = py_touserdata(argv);
+PK_DEFINE_NEXT_WRAPPER(generator)
+
+int generator__iternext(py_Ref self) {
+    Generator* ud = py_touserdata(self);
     py_StackRef p0 = py_peek(0);
     VM* vm = pk_current_vm;
-    if(ud->state == 2) return StopIteration();
+    if(ud->state == 2) {
+        py_newnil(py_retval());
+        return 0;
+    }
 
     // reset frame->p0
     assert(!ud->frame->is_locals_special);
@@ -35,7 +41,7 @@ bool generator__next__(int argc, py_Ref argv) {
     ud->frame->locals = ud->frame->p0 + locals_offset;
     
     // restore the context
-    py_Ref backup = py_getslot(argv, 0);
+    py_Ref backup = py_getslot(self, 0);
     int length = py_list_len(backup);
     py_TValue* p = py_list_data(backup);
     for(int i = 0; i < length; i++)
@@ -51,10 +57,19 @@ bool generator__next__(int argc, py_Ref argv) {
     if(res == RES_ERROR) {
         ud->state = 2;  // end this generator immediately on error
         if(py_matchexc(tp_StopIteration)) {
+            // PEP 479: a `StopIteration` escaping the body must not be mistaken
+            // for the generator finishing normally
+            py_TValue stop_iter = *py_retval();  // stashed there by py_matchexc
             py_clearexc(p0);
-            return true;
+            // root it on the stack, `RuntimeError` below allocates
+            py_StackRef inner = py_pushtmp();
+            *inner = stop_iter;
+            RuntimeError("generator raised StopIteration");
+            BaseException* exc = py_touserdata(&vm->unhandled_exc);
+            exc->inner_exc = *inner;
+            py_pop();
         }
-        return false;
+        return -1;
     }
 
     if(res == RES_YIELD) {
@@ -67,14 +82,13 @@ bool generator__next__(int argc, py_Ref argv) {
         vm->top_frame = vm->top_frame->f_back;
         vm->recursion_depth--;
         ud->state = 1;
-        return true;
+        return 1;
     } else {
         assert(res == RES_RETURN);
         ud->state = 2;
-        // raise StopIteration(<retval>)
-        bool ok = py_tpcall(tp_StopIteration, 1, py_retval());
-        if(!ok) return false;
-        return py_raise(py_retval());
+        // `py_retval()` already holds the return value, which the caller turns
+        // into `StopIteration(<retval>)` if it needs a real exception
+        return 0;
     }
 }
 

src/modules/array2d.c

@@ -1,5 +1,6 @@
 #include "pocketpy/interpreter/array2d.h"
 #include "pocketpy/interpreter/vm.h"
+#include "pocketpy/interpreter/bindings.h"
 #include "pocketpy/pocketpy.h"
 #include <limits.h>
 
@@ -884,10 +885,14 @@ static void register_array2d_like(py_Ref mod) {
     }
 }
 
-bool array2d_like_iterator__next__(int argc, py_Ref argv) {
-    PY_CHECK_ARGC(1);
-    c11_array2d_like_iterator* self = py_touserdata(argv);
-    if(self->j >= self->array->n_rows) return StopIteration();
+PK_DEFINE_NEXT_WRAPPER(array2d_like_iterator)
+
+int array2d_like_iterator__iternext(py_Ref self_) {
+    c11_array2d_like_iterator* self = py_touserdata(self_);
+    if(self->j >= self->array->n_rows) {
+        py_newnil(py_retval());
+        return 0;
+    }
     py_TValue* data = py_newtuple(py_retval(), 2);
     py_newvec2i(&data[0],
                 (c11_vec2i){
@@ -899,7 +904,7 @@ bool array2d_like_iterator__next__(int argc, py_Ref argv) {
         self->i = 0;
         self->j++;
     }
-    return true;
+    return 1;
 }
 
 static void register_array2d_like_iterator(py_Ref mod) {

src/modules/builtins.c

@@ -7,6 +7,7 @@
 #include "pocketpy/objects/object.h"
 #include "pocketpy/common/sstream.h"
 #include "pocketpy/interpreter/vm.h"
+#include "pocketpy/interpreter/bindings.h"
 #include "pocketpy/common/_generated.h"
 
 
@@ -123,8 +124,8 @@ static bool builtins_next(int argc, py_Ref argv) {
     if(res == -1) return false;
     if(res) return true;
     if(argc == 1) {
-        // StopIteration stored in py_retval()
-        return py_raise(py_retval());
+        // py_retval() holds the StopIteration value, or nil if there is none
+        return pk__raise_stopiteration();
     } else {
         py_assign(py_retval(), py_arg(1));
         return true;

src/public/PyDict.c

@@ -4,6 +4,7 @@
 #include "pocketpy/common/sstream.h"
 #include "pocketpy/interpreter/types.h"
 #include "pocketpy/interpreter/vm.h"
+#include "pocketpy/interpreter/bindings.h"
 
 typedef struct {
     Dict* dict;  // weakref for slot 0
@@ -644,30 +645,35 @@ py_Type pk_dict__register() {
 }
 
 //////////////////////////
-bool dict_items__next__(int argc, py_Ref argv) {
-    PY_CHECK_ARGC(1);
-    DictIterator* iter = py_touserdata(py_arg(0));
-    if(DictIterator__modified(iter)) return RuntimeError("dictionary modified during iteration");
+PK_DEFINE_NEXT_WRAPPER(dict_items)
+
+int dict_items__iternext(py_Ref self) {
+    DictIterator* iter = py_touserdata(self);
+    if(DictIterator__modified(iter)) {
+        RuntimeError("dictionary modified during iteration");
+        return -1;
+    }
     DictEntry* entry = (DictIterator__next(iter));
     if(entry) {
         switch(iter->mode) {
             case 0:  // keys
                 py_assign(py_retval(), &entry->key);
-                return true;
+                return 1;
             case 1:  // values
                 py_assign(py_retval(), &entry->val);
-                return true;
+                return 1;
             case 2:  // items
             {
                 py_Ref p = py_newtuple(py_retval(), 2);
                 p[0] = entry->key;
                 p[1] = entry->val;
-                return true;
+                return 1;
             }
             default: c11__unreachable();
         }
     }
-    return StopIteration();
+    py_newnil(py_retval());
+    return 0;
 }
 
 bool dict_items__len__(int argc, py_Ref argv) {

src/public/PyException.c

@@ -5,6 +5,7 @@
 #include "pocketpy/interpreter/vm.h"
 #include "pocketpy/common/sstream.h"
 #include "pocketpy/objects/exception.h"
+#include "pocketpy/interpreter/bindings.h"
 
 void py_BaseException__stpush(py_Frame* frame,
                               py_Ref self,
@@ -311,3 +312,10 @@ bool StopIteration() {
     if(!ok) return false;
     return py_raise(py_retval());
 }
+
+bool pk__raise_stopiteration() {
+    if(py_isnil(py_retval())) return StopIteration();
+    // carry the value, e.g. `StopIteration(<generator return value>)`
+    if(!py_tpcall(tp_StopIteration, 1, py_retval())) return false;
+    return py_raise(py_retval());
+}

src/public/PythonOps.c

@@ -2,6 +2,7 @@
 #include "pocketpy/interpreter/bindings.h"
 #include "pocketpy/interpreter/vm.h"
 #include "pocketpy/objects/base.h"
+#include "pocketpy/objects/exception.h"
 #include "pocketpy/pocketpy.h"
 
 bool py_binaryadd(py_Ref lhs, py_Ref rhs) { return py_binaryop(lhs, rhs, __add__, __radd__); }
@@ -142,43 +143,31 @@ bool py_iter(py_Ref val) {
 }
 
 int py_next(py_Ref val) {
-    VM* vm = pk_current_vm;
-
+    // builtin iterators signal exhaustion without constructing a `StopIteration`
     switch(val->type) {
-        case tp_generator:
-            if(generator__next__(1, val)) return 1;
-            break;
-        case tp_array2d_like_iterator:
-            if(array2d_like_iterator__next__(1, val)) return 1;
-            break;
-        case tp_list_iterator:
-            if(list_iterator__next__(1, val)) return 1;
-            break;
-        case tp_tuple_iterator:
-            if(tuple_iterator__next__(1, val)) return 1;
-            break;
-        case tp_dict_iterator:
-            if(dict_items__next__(1, val)) return 1;
-            break;
-        case tp_range_iterator:
-            if(range_iterator__next__(1, val)) return 1;
-            break;
-        case tp_str_iterator:
-            if(str_iterator__next__(1, val)) return 1;
-            break;
-        default: {
-            py_Ref tmp = py_tpfindmagic(val->type, __next__);
-            if(!tmp) {
-                TypeError("'%t' object is not an iterator", val->type);
-                return -1;
-            }
-            if(py_call(tmp, 1, val)) return 1;
-            break;
-        }
+        case tp_generator: return generator__iternext(val);
+        case tp_array2d_like_iterator: return array2d_like_iterator__iternext(val);
+        case tp_list_iterator: return list_iterator__iternext(val);
+        case tp_tuple_iterator: return tuple_iterator__iternext(val);
+        case tp_dict_iterator: return dict_items__iternext(val);
+        case tp_range_iterator: return range_iterator__iternext(val);
+        case tp_str_iterator: return str_iterator__iternext(val);
+        default: break;
+    }
+
+    VM* vm = pk_current_vm;
+    py_Ref tmp = py_tpfindmagic(val->type, __next__);
+    if(!tmp) {
+        TypeError("'%t' object is not an iterator", val->type);
+        return -1;
     }
+    if(py_call(tmp, 1, val)) return 1;
     if(vm->unhandled_exc.type == tp_StopIteration) {
-        vm->last_retval = vm->unhandled_exc;
+        // unwrap the value so callers never have to touch the exception object
+        BaseException* ud = py_touserdata(&vm->unhandled_exc);
+        py_TValue value = ud->args;
         py_clearexc(NULL);
+        *py_retval() = value;
         return 0;
     }
     return -1;

tests/290_iter.py

@@ -50,3 +50,85 @@ try:
 except StopIteration:
     pass
 
+
+# --- StopIteration carries the right value ------------------------------
+# `py_next` reports exhaustion without building a StopIteration object, so the
+# object has to be reconstructed faithfully wherever one is actually observable.
+
+it = iter([1])
+assert next(it) == 1
+try:
+    next(it)
+    exit(1)
+except StopIteration as e:
+    assert e.args == ()
+    assert e.value is None
+    assert repr(e) == 'StopIteration()'
+
+assert next(iter([]), 'dflt') == 'dflt'
+
+def gen_with_return():
+    yield 1
+    return 42
+
+it = iter(gen_with_return())
+assert next(it) == 1
+try:
+    next(it)
+    exit(1)
+except StopIteration as e:
+    assert e.args == (42,)
+    assert e.value == 42
+
+def gen_bare_return():
+    yield 1
+
+it = iter(gen_bare_return())
+assert next(it) == 1
+try:
+    next(it)
+    exit(1)
+except StopIteration as e:
+    assert e.value is None
+
+# `yield from` reads the value out of the exhausted sub-iterator
+def outer_with_return():
+    got = yield from gen_with_return()
+    yield got
+assert list(outer_with_return()) == [1, 42]
+
+def outer_bare_return():
+    got = yield from gen_bare_return()
+    yield got
+assert list(outer_bare_return()) == [1, None]
+
+class RaisesWithValue:
+    def __iter__(self):
+        return self
+    def __next__(self):
+        raise StopIteration('V')
+
+def outer_user_iter():
+    got = yield from RaisesWithValue()
+    yield got
+assert list(outer_user_iter()) == ['V']
+
+# --- every builtin iterator still terminates ---------------------------
+assert list(iter([1, 2])) == [1, 2]
+assert list(iter((1, 2))) == [1, 2]
+assert list(range(3)) == [0, 1, 2]
+assert list('ab') == ['a', 'b']
+assert ''.join(iter(['a', 'b'])) == 'ab'
+
+d = {'x': 1, 'y': 2}
+assert sorted(d.keys()) == ['x', 'y']
+assert sorted(d.values()) == [1, 2]
+assert sorted(d.items()) == [('x', 1), ('y', 2)]
+
+# a dict mutated mid-iteration must still be reported as an error
+try:
+    for k in d:
+        d['z'] = 3
+    exit(1)
+except RuntimeError:
+    pass

tests/510_yield.py

@@ -127,4 +127,60 @@ def f():
     a = yield from g()
     yield a
 
-assert list(f()) == [1, 2, 3]
+assert list(f()) == [1, 2, 3]
+# --- PEP 479: a StopIteration escaping a generator body becomes RuntimeError ---
+# NOTE: the builtin `iter` is shadowed above, so build iterators via generators
+def _exhausted():
+    return
+    yield
+
+def raises_stop_iteration():
+    yield 1
+    raise StopIteration
+
+try:
+    list(raises_stop_iteration())
+    exit(1)
+except RuntimeError as e:
+    assert str(e) == 'generator raised StopIteration', str(e)
+
+it = raises_stop_iteration()
+assert next(it) == 1
+try:
+    next(it)
+    exit(1)
+except RuntimeError:
+    pass
+
+# the same applies when it comes from an exhausted inner iterator
+def drains_inner():
+    inner = _exhausted()
+    yield 1
+    next(inner)
+
+try:
+    list(drains_inner())
+    exit(1)
+except RuntimeError:
+    pass
+
+# a generator that catches it itself is unaffected
+def catches_it():
+    try:
+        next(_exhausted())
+        exit(1)
+    except StopIteration:
+        yield 'caught'
+
+assert list(catches_it()) == ['caught']
+
+# `yield from` over a sub-generator that finishes normally is unaffected
+def sub_with_return():
+    yield 'a'
+    return 'R'
+
+def delegates():
+    got = yield from sub_with_return()
+    yield got
+
+assert list(delegates()) == ['a', 'R']