Merge pull request #525 from toshok/fewer-stopiter-exceptions

Fewer StopIteration exceptions

from_cpython/Objects/iterobject.c

@@ -9,6 +9,8 @@ typedef struct {
     PyObject_HEAD
     PyObject *it_callable; /* Set to NULL when iterator is exhausted */
     PyObject *it_sentinel; /* Set to NULL when iterator is exhausted */
+    // Pyston changes:
+    PyObject *it_nextvalue; /* Set to non-null when iterator is advanced in __hasnext__ */
 } calliterobject;
 
 PyObject *
@@ -22,6 +24,8 @@ PyCallIter_New(PyObject *callable, PyObject *sentinel)
     it->it_callable = callable;
     Py_INCREF(sentinel);
     it->it_sentinel = sentinel;
+    // Pyston change: initialize nextvalue
+    it->it_nextvalue = NULL;
     _PyObject_GC_TRACK(it);
     return (PyObject *)it;
 }
@@ -31,6 +35,8 @@ calliter_dealloc(calliterobject *it)
     _PyObject_GC_UNTRACK(it);
     Py_XDECREF(it->it_callable);
     Py_XDECREF(it->it_sentinel);
+    // Pyston change: xdecref nextvalue
+    Py_XDECREF(it->it_nextvalue);
     PyObject_GC_Del(it);
 }
 
@@ -39,11 +45,15 @@ calliter_traverse(calliterobject *it, visitproc visit, void *arg)
 {
     Py_VISIT(it->it_callable);
     Py_VISIT(it->it_sentinel);
+    // Pyston change: visit nextvalue
+    Py_VISIT(it->it_nextvalue);
     return 0;
 }
 
+// Pyston change: extract most of the body of calliter_iternext here
+// so we can use it from both calliter_iternext and calliter_hasnext
 static PyObject *
-calliter_iternext(calliterobject *it)
+calliter_next(calliterobject *it)
 {
     if (it->it_callable != NULL) {
         PyObject *args = PyTuple_New(0);
@@ -69,11 +79,43 @@ calliter_iternext(calliterobject *it)
             PyErr_Clear();
             Py_CLEAR(it->it_callable);
             Py_CLEAR(it->it_sentinel);
+            Py_CLEAR(it->it_nextvalue);
         }
     }
     return NULL;
 }
 
+static PyObject *
+calliter_iternext(calliterobject *it)
+{
+    // Pyston change: for __hasnext__ based iteration, return the cached next value
+    if (it->it_nextvalue != NULL) {
+        PyObject* rv = it->it_nextvalue;
+        Py_CLEAR(it->it_nextvalue);
+        return rv;
+    }
+
+    return calliter_next(it);
+}
+
+// Pyston addition: __hasnext__ based iteration
+static int
+calliter_hasnext(calliterobject *it)
+{
+    if (!it->it_nextvalue) {
+        it->it_nextvalue = calliter_next(it);
+    }
+    return it->it_nextvalue != NULL;
+}
+
+// Pyston change: give iter objects a __hasnext__ method
+void
+PyCallIter_AddHasNext()
+{
+    PyCallIter_Type._tpp_hasnext = calliter_hasnext;
+}
+
+
 PyTypeObject PyCallIter_Type = {
     // Pyston change:
     PyVarObject_HEAD_INIT(NULL /* &PyType_Type */, 0)

src/capi/types.h

@@ -152,7 +152,7 @@ public:
         if (flags == PyWrapperFlag_KEYWORDS) {
             wrapperfunc_kwds wk = (wrapperfunc_kwds)wrapper;
             rtn = (*wk)(self->obj, args, self->descr->wrapped, kwds);
-        } else if (flags == 0) {
+        } else if (flags == PyWrapperFlag_PYSTON || flags == 0) {
             rtn = (*wrapper)(self->obj, args, self->descr->wrapped);
         } else {
             RELEASE_ASSERT(0, "%d", flags);

src/runtime/generator.cpp

@@ -123,17 +123,15 @@ Box* generatorIter(Box* s) {
     return s;
 }
 
-Box* generatorSend(Box* s, Box* v) {
-    assert(s->cls == generator_cls);
-    BoxedGenerator* self = static_cast<BoxedGenerator*>(s);
-
+// called from both generatorHasNext and generatorSend/generatorNext (but only if generatorHasNext hasn't been called)
+static void generatorSendInternal(BoxedGenerator* self, Box* v) {
     if (self->running)
         raiseExcHelper(ValueError, "generator already executing");
 
     // check if the generator already exited
     if (self->entryExited) {
         freeGeneratorStack(self);
-        raiseExcHelper(StopIteration, "");
+        return;
     }
 
     self->returnValue = v;
@@ -161,9 +159,23 @@ Box* generatorSend(Box* s, Box* v) {
         raiseRaw(self->exception);
     }
 
-    // throw StopIteration if the generator exited
     if (self->entryExited) {
         freeGeneratorStack(self);
+        return;
+    }
+}
+
+Box* generatorSend(Box* s, Box* v) {
+    assert(s->cls == generator_cls);
+    BoxedGenerator* self = static_cast<BoxedGenerator*>(s);
+
+    if (self->iterated_from__hasnext__)
+        Py_FatalError(".throw called on generator last advanced with __hasnext__");
+
+    generatorSendInternal(self, v);
+
+    // throw StopIteration if the generator exited
+    if (self->entryExited) {
         raiseExcHelper(StopIteration, "");
     }
 
@@ -173,6 +185,10 @@ Box* generatorSend(Box* s, Box* v) {
 Box* generatorThrow(Box* s, BoxedClass* exc_cls, Box* exc_val = nullptr, Box** args = nullptr) {
     assert(s->cls == generator_cls);
     BoxedGenerator* self = static_cast<BoxedGenerator*>(s);
+
+    if (self->iterated_from__hasnext__)
+        Py_FatalError(".throw called on generator last advanced with __hasnext__");
+
     Box* exc_tb = args ? nullptr : args[0];
     if (!exc_val)
         exc_val = None;
@@ -194,9 +210,34 @@ Box* generatorClose(Box* s) {
 }
 
 Box* generatorNext(Box* s) {
+    assert(s->cls == generator_cls);
+    BoxedGenerator* self = static_cast<BoxedGenerator*>(s);
+
+    if (self->iterated_from__hasnext__) {
+        self->iterated_from__hasnext__ = false;
+        return self->returnValue;
+    }
+
     return generatorSend(s, None);
 }
 
+i1 generatorHasnextUnboxed(Box* s) {
+    assert(s->cls == generator_cls);
+    BoxedGenerator* self = static_cast<BoxedGenerator*>(s);
+
+    if (!self->iterated_from__hasnext__) {
+        generatorSendInternal(self, None);
+        self->iterated_from__hasnext__ = true;
+    }
+
+    return !self->entryExited;
+}
+
+Box* generatorHasnext(Box* s) {
+    return boxBool(generatorHasnextUnboxed(s));
+}
+
+
 extern "C" Box* yield(BoxedGenerator* obj, Box* value) {
     assert(obj->cls == generator_cls);
     BoxedGenerator* self = static_cast<BoxedGenerator*>(obj);
@@ -368,6 +409,11 @@ void setupGenerator() {
 
     generator_cls->giveAttr("close", new BoxedFunction(boxRTFunction((void*)generatorClose, UNKNOWN, 1)));
     generator_cls->giveAttr("next", new BoxedFunction(boxRTFunction((void*)generatorNext, UNKNOWN, 1)));
+
+    CLFunction* hasnext = boxRTFunction((void*)generatorHasnextUnboxed, BOOL, 1);
+    addRTFunction(hasnext, (void*)generatorHasnext, BOXED_BOOL);
+    generator_cls->giveAttr("__hasnext__", new BoxedFunction(hasnext));
+
     generator_cls->giveAttr("send", new BoxedFunction(boxRTFunction((void*)generatorSend, UNKNOWN, 2)));
     auto gthrow = new BoxedFunction(boxRTFunction((void*)generatorThrow, UNKNOWN, 4, 2, false, false), { NULL, NULL });
     generator_cls->giveAttr("throw", gthrow);

src/runtime/types.cpp

@@ -2058,6 +2058,8 @@ inline void initUserAttrs(Box* obj, BoxedClass* cls) {
     }
 }
 
+extern "C" void PyCallIter_AddHasNext();
+
 extern "C" PyObject* PyObject_Init(PyObject* op, PyTypeObject* tp) noexcept {
     assert(op);
     assert(tp);
@@ -2512,6 +2514,7 @@ void setupRuntime() {
 
     PyType_Ready(&PyByteArrayIter_Type);
     PyType_Ready(&PyCapsule_Type);
+    PyCallIter_AddHasNext();
     PyType_Ready(&PyCallIter_Type);
     PyType_Ready(&PyCObject_Type);
 

src/runtime/types.h

@@ -834,6 +834,7 @@ public:
     bool entryExited;
     bool running;
     Box* returnValue;
+    bool iterated_from__hasnext__;
     ExcInfo exception;
 
     struct Context* context, *returnContext;