Support StopIteration-based for loop iteration

At runtime, we can detect if something supports the new Pyston iteration
protocol (__hasnext__) or not.

Statically, when we lower a for loop, we assume that the iterator supports
the new protocol, and we can't check it

This commit adds a check to getiter() which wraps a Python-style iterator
with an IterWrapper which adds the __hasnext__ method.

src/codegen/irgen/irgenerator.cpp

@@ -465,6 +465,9 @@ private:
                 // TODO if this is a type that has an __iter__, we could do way better than this, both in terms of
                 // function call overhead and resulting type information, if we went with that instead of the generic
                 // version.
+                // (ie we can inline getiter here, whether mechanically with LLVM [would require adding more
+                // optimization passes to make it fast] or by-hand)
+                //
                 // TODO Move this behavior into to the type-specific section (compvars.cpp)?
                 emitter.getBuilder();
                 assert(node->args.size() == 1);

src/runtime/iterobject.cpp

@@ -30,6 +30,7 @@
 namespace pyston {
 
 BoxedClass* seqiter_cls;
+BoxedClass* iterwrapper_cls;
 
 Box* seqiterHasnext(Box* s) {
     RELEASE_ASSERT(s->cls == seqiter_cls, "");
@@ -56,6 +57,45 @@ Box* seqiterNext(Box* s) {
     return r;
 }
 
+static void iterwrapperGCVisit(GCVisitor* v, Box* b) {
+    assert(b->cls == iterwrapper_cls);
+    boxGCHandler(v, b);
+
+    BoxedIterWrapper* iw = static_cast<BoxedIterWrapper*>(b);
+    if (iw->next)
+        v->visit(iw->next);
+}
+
+Box* iterwrapperHasnext(Box* s) {
+    RELEASE_ASSERT(s->cls == iterwrapper_cls, "");
+    BoxedIterWrapper* self = static_cast<BoxedIterWrapper*>(s);
+
+    static const std::string next_str("next");
+    Box* next;
+    try {
+        next = callattr(self->iter, &next_str, CallattrFlags({.cls_only = true, .null_on_nonexistent = false }),
+                        ArgPassSpec(0), NULL, NULL, NULL, NULL, NULL);
+    } catch (Box* b) {
+        if (isSubclass(b->cls, StopIteration)) {
+            self->next = NULL;
+            return False;
+        }
+        throw;
+    }
+    self->next = next;
+    return True;
+}
+
+Box* iterwrapperNext(Box* s) {
+    RELEASE_ASSERT(s->cls == iterwrapper_cls, "");
+    BoxedIterWrapper* self = static_cast<BoxedIterWrapper*>(s);
+
+    RELEASE_ASSERT(self->next, "");
+    Box* r = self->next;
+    self->next = NULL;
+    return r;
+}
+
 void setupIter() {
     seqiter_cls = new BoxedHeapClass(object_cls, NULL, 0, sizeof(BoxedSeqIter), false);
     seqiter_cls->giveAttr("__name__", boxStrConstant("iterator"));
@@ -64,5 +104,14 @@ void setupIter() {
     seqiter_cls->giveAttr("__hasnext__", new BoxedFunction(boxRTFunction((void*)seqiterHasnext, BOXED_BOOL, 1)));
 
     seqiter_cls->freeze();
+
+    iterwrapper_cls = new BoxedHeapClass(object_cls, iterwrapperGCVisit, 0, sizeof(BoxedIterWrapper), false);
+    iterwrapper_cls->giveAttr("__name__", boxStrConstant("iterwrapper"));
+
+    iterwrapper_cls->giveAttr("next", new BoxedFunction(boxRTFunction((void*)iterwrapperNext, UNKNOWN, 1)));
+    iterwrapper_cls->giveAttr("__hasnext__",
+                              new BoxedFunction(boxRTFunction((void*)iterwrapperHasnext, BOXED_BOOL, 1)));
+
+    iterwrapper_cls->freeze();
 }
 }

src/runtime/iterobject.h

@@ -23,6 +23,9 @@
 namespace pyston {
 
 extern BoxedClass* seqiter_cls;
+
+// Analogue of CPython's PySeqIter: wraps an object that has a __getitem__
+// and uses that to iterate.
 class BoxedSeqIter : public Box {
 public:
     Box* b;
@@ -34,6 +37,19 @@ public:
     DEFAULT_CLASS(seqiter_cls);
 };
 
+extern BoxedClass* iterwrapper_cls;
+// Pyston wrapper that wraps CPython-style iterators (next() which throws StopException)
+// and converts it to Pyston-style (__hasnext__)
+class BoxedIterWrapper : public Box {
+public:
+    Box* iter;
+    Box* next;
+
+    BoxedIterWrapper(Box* iter) : iter(iter) {}
+
+    DEFAULT_CLASS(iterwrapper_cls);
+};
+
 void setupIter();
 }
 

src/runtime/objmodel.cpp

@@ -3353,8 +3353,12 @@ extern "C" Box* getiter(Box* o) {
     // TODO add rewriting to this?  probably want to try to avoid this path though
     static const std::string iter_str("__iter__");
     Box* r = callattrInternal0(o, &iter_str, LookupScope::CLASS_ONLY, NULL, ArgPassSpec(0));
-    if (r)
+    if (r) {
+        static const std::string hasnext_str("__hasnext__");
+        if (typeLookup(r->cls, hasnext_str, NULL) == NULL)
+            return new BoxedIterWrapper(r);
         return r;
+    }
 
     static const std::string getitem_str("__getitem__");
     if (typeLookup(o->cls, getitem_str, NULL)) {

test/tests/for_iter.py

-# expected: fail
-# - real iteration protocol is unsupported (no exceptions yet)
-
 class C(object):
     def __iter__(self):
         print "orig iter"
@@ -36,3 +33,19 @@ c.next = newnext
 
 for i in c: # should hit the old next
     print i
+
+class C(object):
+    def __init__(self):
+        self.n = 0
+
+    def __iter__(self):
+        return self
+
+    def next(self):
+        if self.n < 10:
+            self.n += 1
+            return self.n * self.n
+        raise StopIteration()
+
+for i in C():
+    print i