Merge branch 'iteration'

src/codegen/ast_interpreter.cpp

@@ -530,7 +530,7 @@ Value ASTInterpreter::visit_langPrimitive(AST_LangPrimitive* node) {
     Value v;
     if (node->opcode == AST_LangPrimitive::GET_ITER) {
         assert(node->args.size() == 1);
-        v = getiter(visit_expr(node->args[0]).o);
+        v = getPystonIter(visit_expr(node->args[0]).o);
     } else if (node->opcode == AST_LangPrimitive::IMPORT_FROM) {
         assert(node->args.size() == 2);
         assert(node->args[0]->type == AST_TYPE::Name);

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 getPystonIter 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);
@@ -473,7 +476,7 @@ private:
                 ConcreteCompilerVariable* converted_obj = obj->makeConverted(emitter, obj->getBoxType());
                 obj->decvref(emitter);
 
-                llvm::Value* v = emitter.createCall(unw_info, g.funcs.getiter, { converted_obj->getValue() });
+                llvm::Value* v = emitter.createCall(unw_info, g.funcs.getPystonIter, { converted_obj->getValue() });
                 assert(v->getType() == g.llvm_value_type_ptr);
 
                 return new ConcreteCompilerVariable(UNKNOWN, v, true);

src/codegen/runtime_hooks.cpp

@@ -197,7 +197,7 @@ void initGlobalFuncs(GlobalState& g) {
     GET(str);
     GET(isinstance);
     GET(yield);
-    GET(getiter);
+    GET(getPystonIter);
 
     GET(unpackIntoArray);
     GET(raiseAttributeError);

src/codegen/runtime_hooks.h

@@ -37,7 +37,7 @@ struct GlobalFuncs {
         *createUserClass, *createClosure, *createGenerator, *createLong, *createSet, *createPureImaginary;
     llvm::Value* getattr, *setattr, *delattr, *delitem, *delGlobal, *nonzero, *binop, *compare, *augbinop, *unboxedLen,
         *getitem, *getclsattr, *getGlobal, *setitem, *unaryop, *import, *importFrom, *importStar, *repr, *str,
-        *isinstance, *yield, *getiter;
+        *isinstance, *yield, *getPystonIter;
 
     llvm::Value* unpackIntoArray, *raiseAttributeError, *raiseAttributeErrorStr, *raiseNotIterableError,
         *assertNameDefined, *assertFail;

src/gc/collector.cpp

@@ -138,8 +138,8 @@ void GCVisitor::visit(void* p) {
 }
 
 void GCVisitor::visitRange(void* const* start, void* const* end) {
-    ASSERT((char*)end - (char*)start <= 1000000000, "Asked to scan %.1fGB -- a bug?",
-           ((char*)end - (char*)start) * 1.0 / (1 << 30));
+    ASSERT((const char*)end - (const char*)start <= 1000000000, "Asked to scan %.1fGB -- a bug?",
+           ((const char*)end - (const char*)start) * 1.0 / (1 << 30));
 
     assert((uintptr_t)start % sizeof(void*) == 0);
     assert((uintptr_t)end % sizeof(void*) == 0);
@@ -158,8 +158,8 @@ void GCVisitor::visitPotential(void* p) {
 }
 
 void GCVisitor::visitPotentialRange(void* const* start, void* const* end) {
-    ASSERT((char*)end - (char*)start <= 1000000000, "Asked to scan %.1fGB -- a bug?",
-           ((char*)end - (char*)start) * 1.0 / (1 << 30));
+    ASSERT((const char*)end - (const char*)start <= 1000000000, "Asked to scan %.1fGB -- a bug?",
+           ((const char*)end - (const char*)start) * 1.0 / (1 << 30));
 
     assert((uintptr_t)start % sizeof(void*) == 0);
     assert((uintptr_t)end % sizeof(void*) == 0);

src/runtime/inline/link_forcer.cpp

@@ -90,7 +90,7 @@ void force() {
     FORCE(str);
     FORCE(isinstance);
     FORCE(yield);
-    FORCE(getiter);
+    FORCE(getPystonIter);
 
     FORCE(unpackIntoArray);
     FORCE(raiseAttributeError);

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

@@ -3349,7 +3349,15 @@ extern "C" void delattr(Box* obj, const char* attr) {
     delattr_internal(obj, attr, true, NULL);
 }
 
-extern "C" Box* getiter(Box* o) {
+extern "C" Box* getPystonIter(Box* o) {
+    Box* r = getiter(o);
+    static const std::string hasnext_str("__hasnext__");
+    if (typeLookup(r->cls, hasnext_str, NULL) == NULL)
+        return new BoxedIterWrapper(r);
+    return r;
+}
+
+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));
@@ -3365,6 +3373,7 @@ extern "C" Box* getiter(Box* o) {
 }
 
 llvm::iterator_range<BoxIterator> Box::pyElements() {
+    // TODO: this should probably call getPystonIter
     Box* iter = getiter(this);
     assert(iter);
     return llvm::iterator_range<BoxIterator>(++BoxIterator(iter), BoxIterator(nullptr));

src/runtime/objmodel.h

@@ -84,7 +84,9 @@ extern "C" void assertFail(BoxedModule* inModule, Box* msg);
 extern "C" bool isInstance(Box* obj, BoxedClass* parent);
 extern "C" bool isSubclass(BoxedClass* child, BoxedClass* parent);
 extern "C" BoxedClosure* createClosure(BoxedClosure* parent_closure);
-extern "C" Box* getiter(Box* o);
+
+Box* getiter(Box* o);
+extern "C" Box* getPystonIter(Box* o);
 
 extern "C" void dump(void* p);
 

test/tests/builtins.py

@@ -68,3 +68,13 @@ try:
     print hex([])
 except TypeError, e:
     print e
+
+class Iterable(object):
+    def __iter__(self):
+        return self
+    def next(self):
+        return 1
+
+i = Iterable()
+it = iter(i)
+print it is i

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