Merge pull request #888 from undingen/wrapper4

add tp_nextiter implementation for our iterators

src/codegen/irgen/irgenerator.cpp

@@ -534,7 +534,7 @@ bool isIsDefinedName(llvm::StringRef name) {
 
 InternedString getIsDefinedName(InternedString name, InternedStringPool& interned_strings) {
     // TODO could cache this
-    return interned_strings.get("!is_defined_" + std::string(name.s()));
+    return interned_strings.get(("!is_defined_" + name.s()).str());
 }
 
 class IRGeneratorImpl : public IRGenerator {

src/core/util.cpp

@@ -96,18 +96,6 @@ Timer::~Timer() {
 
 #endif // !DISABLE_TIMERS
 
-bool startswith(const std::string& s, const std::string& pattern) {
-    if (pattern.size() > s.size())
-        return false;
-    return s.compare(0, pattern.size(), pattern) == 0;
-}
-
-bool endswith(const std::string& s, const std::string& pattern) {
-    if (pattern.size() > s.size())
-        return false;
-    return s.compare(s.size() - pattern.size(), pattern.size(), pattern) == 0;
-}
-
 void removeDirectoryIfExists(const std::string& path) {
     llvm_error_code code;
 

src/core/util.h

@@ -82,8 +82,13 @@ public:
 
 #endif // #else DISABLE_TIMERS
 
-bool startswith(const std::string& s, const std::string& pattern);
-bool endswith(const std::string& s, const std::string& pattern);
+inline bool startswith(llvm::StringRef s, llvm::StringRef pattern) {
+    return s.startswith(pattern);
+}
+
+inline bool endswith(llvm::StringRef s, llvm::StringRef pattern) {
+    return s.endswith(pattern);
+}
 
 void removeDirectoryIfExists(const std::string& path);
 

src/runtime/dict.cpp

@@ -848,8 +848,8 @@ void setupDict() {
     dict_iterator_cls->giveAttr(
         "__iter__", new BoxedFunction(boxRTFunction((void*)dictIterIter, typeFromClass(dict_iterator_cls), 1)));
     dict_iterator_cls->giveAttr("next", new BoxedFunction(boxRTFunction((void*)dictIterNext, UNKNOWN, 1)));
-
     dict_iterator_cls->freeze();
+    dict_iterator_cls->tp_iternext = dictiter_next;
 
     // Manually set some tp_* slots *after* calling freeze() -> fixup_slot_dispatchers().
     // fixup_slot_dispatchers will insert a wrapper like slot_tp_init into tp_init, which calls the python-level

src/runtime/dict.h

@@ -49,6 +49,7 @@ Box* dictIterItems(Box* self);
 Box* dictIterIter(Box* self);
 Box* dictIterHasnext(Box* self);
 i1 dictIterHasnextUnboxed(Box* self);
+Box* dictiter_next(Box* self) noexcept;
 Box* dictIterNext(Box* self);
 
 class BoxedDictView : public Box {

src/runtime/inline/dict.cpp

@@ -56,12 +56,12 @@ Box* dictIterHasnext(Box* s) {
     return boxBool(dictIterHasnextUnboxed(s));
 }
 
-Box* dictIterNext(Box* s) {
+Box* dictiter_next(Box* s) noexcept {
     assert(s->cls == dict_iterator_cls);
     BoxedDictIterator* self = static_cast<BoxedDictIterator*>(s);
 
     if (self->it == self->itEnd)
-        raiseExcHelper(StopIteration, "");
+        return NULL;
 
     Box* rtn = nullptr;
     if (self->type == BoxedDictIterator::KeyIterator) {
@@ -75,6 +75,13 @@ Box* dictIterNext(Box* s) {
     return rtn;
 }
 
+Box* dictIterNext(Box* s) {
+    auto* rtn = dictiter_next(s);
+    if (!rtn)
+        raiseExcHelper(StopIteration, "");
+    return rtn;
+}
+
 BoxedDictView::BoxedDictView(BoxedDict* d) : d(d) {
 }
 

src/runtime/inline/list.cpp

@@ -65,31 +65,35 @@ i1 listiterHasnextUnboxed(Box* s) {
     return ans;
 }
 
-template <ExceptionStyle S> Box* listiterNext(Box* s) noexcept(S == CAPI) {
+Box* listiter_next(Box* s) noexcept {
     assert(s->cls == list_iterator_cls);
     BoxedListIterator* self = static_cast<BoxedListIterator*>(s);
 
-    if (!self->l) {
-        if (S == CAPI) {
-            PyErr_SetObject(StopIteration, None);
-            return NULL;
-        } else
-            raiseExcHelper(StopIteration, "");
-    }
+    if (!self->l)
+        return NULL;
 
     if (!(self->pos >= 0 && self->pos < self->l->size)) {
         self->l = NULL;
+        return NULL;
+    }
+
+    Box* rtn = self->l->elts->elts[self->pos];
+    self->pos++;
+    return rtn;
+}
+
+template <ExceptionStyle S> Box* listiterNext(Box* s) noexcept(S == CAPI) {
+    Box* rtn = listiter_next(s);
+    if (!rtn) {
         if (S == CAPI) {
             PyErr_SetObject(StopIteration, None);
             return NULL;
         } else
             raiseExcHelper(StopIteration, "");
     }
-
-    Box* rtn = self->l->elts->elts[self->pos];
-    self->pos++;
     return rtn;
 }
+
 // force instantiation:
 template Box* listiterNext<CAPI>(Box*);
 template Box* listiterNext<CXX>(Box*);
@@ -114,12 +118,12 @@ i1 listreviterHasnextUnboxed(Box* s) {
     return self->pos >= 0;
 }
 
-Box* listreviterNext(Box* s) {
+Box* listreviter_next(Box* s) noexcept {
     assert(s->cls == list_reverse_iterator_cls);
     BoxedListIterator* self = static_cast<BoxedListIterator*>(s);
 
     if (!(self->pos >= 0 && self->pos < self->l->size)) {
-        raiseExcHelper(StopIteration, "");
+        return NULL;
     }
 
     Box* rtn = self->l->elts->elts[self->pos];
@@ -127,6 +131,13 @@ Box* listreviterNext(Box* s) {
     return rtn;
 }
 
+Box* listreviterNext(Box* s) {
+    Box* rtn = listreviter_next(s);
+    if (!rtn)
+        raiseExcHelper(StopIteration, "");
+    return rtn;
+}
+
 
 const int BoxedList::INITIAL_CAPACITY = 8;
 // TODO the inliner doesn't want to inline these; is there any point to having them in the inline section?

src/runtime/inline/tuple.cpp

@@ -43,16 +43,24 @@ i1 tupleiterHasnextUnboxed(Box* s) {
     return self->pos < self->t->size();
 }
 
-Box* tupleiterNext(Box* s) {
+Box* tupleiter_next(Box* s) noexcept {
     assert(s->cls == tuple_iterator_cls);
     BoxedTupleIterator* self = static_cast<BoxedTupleIterator*>(s);
 
     if (!(self->pos >= 0 && self->pos < self->t->size())) {
-        raiseExcHelper(StopIteration, "");
+        return NULL;
     }
 
     Box* rtn = self->t->elts[self->pos];
     self->pos++;
     return rtn;
 }
+
+Box* tupleiterNext(Box* s) {
+    Box* rtn = tupleiter_next(s);
+    if (!rtn) {
+        raiseExcHelper(StopIteration, "");
+    }
+    return rtn;
+}
 }

src/runtime/inline/xrange.cpp

@@ -102,6 +102,18 @@ public:
         return boxBool(xrangeIteratorHasnextUnboxed(s));
     }
 
+    static Box* xrangeIterator_next(Box* s) noexcept {
+        assert(s->cls == xrange_iterator_cls);
+        BoxedXrangeIterator* self = static_cast<BoxedXrangeIterator*>(s);
+
+        if (!xrangeIteratorHasnextUnboxed(s))
+            return NULL;
+
+        i64 rtn = self->cur;
+        self->cur += self->step;
+        return boxInt(rtn);
+    }
+
     static i64 xrangeIteratorNextUnboxed(Box* s) __attribute__((visibility("default"))) {
         assert(s->cls == xrange_iterator_cls);
         BoxedXrangeIterator* self = static_cast<BoxedXrangeIterator*>(s);
@@ -228,5 +240,6 @@ void setupXrange() {
 
     xrange_iterator_cls->freeze();
     xrange_iterator_cls->tpp_hasnext = BoxedXrangeIterator::xrangeIteratorHasnextUnboxed;
+    xrange_iterator_cls->tp_iternext = BoxedXrangeIterator::xrangeIterator_next;
 }
 }

src/runtime/list.cpp

@@ -1267,12 +1267,13 @@ void setupList() {
     list_iterator_cls->giveAttr(
         "__iter__", new BoxedFunction(boxRTFunction((void*)listIterIter, typeFromClass(list_iterator_cls), 1)));
 
-    CLFunction* listiter_next = boxRTFunction((void*)listiterNext<CXX>, UNKNOWN, 1);
-    addRTFunction(listiter_next, (void*)listiterNext<CAPI>, UNKNOWN, CAPI);
-    list_iterator_cls->giveAttr("next", new BoxedFunction(listiter_next));
+    CLFunction* listiter_next_func = boxRTFunction((void*)listiterNext<CXX>, UNKNOWN, 1);
+    addRTFunction(listiter_next_func, (void*)listiterNext<CAPI>, UNKNOWN, CAPI);
+    list_iterator_cls->giveAttr("next", new BoxedFunction(listiter_next_func));
 
     list_iterator_cls->freeze();
     list_iterator_cls->tpp_hasnext = listiterHasnextUnboxed;
+    list_iterator_cls->tp_iternext = listiter_next;
 
     list_reverse_iterator_cls->giveAttr("__name__", boxString("listreverseiterator"));
 
@@ -1284,6 +1285,7 @@ void setupList() {
     list_reverse_iterator_cls->giveAttr("next", new BoxedFunction(boxRTFunction((void*)listreviterNext, UNKNOWN, 1)));
 
     list_reverse_iterator_cls->freeze();
+    list_reverse_iterator_cls->tp_iternext = listreviter_next;
 }
 
 void teardownList() {

src/runtime/list.h

@@ -38,10 +38,12 @@ Box* listIterIter(Box* self);
 Box* listiterHasnext(Box* self);
 i1 listiterHasnextUnboxed(Box* self);
 template <ExceptionStyle S> Box* listiterNext(Box* self) noexcept(S == CAPI);
+Box* listiter_next(Box* s) noexcept;
 Box* listReversed(Box* self);
 Box* listreviterHasnext(Box* self);
 i1 listreviterHasnextUnboxed(Box* self);
 Box* listreviterNext(Box* self);
+Box* listreviter_next(Box* s) noexcept;
 void listSort(BoxedList* self, Box* cmp, Box* key, Box* reverse);
 extern "C" Box* listAppend(Box* self, Box* v);
 }

src/runtime/set.cpp

@@ -71,6 +71,14 @@ Box* setiteratorNext(BoxedSetIterator* self) {
     return self->next();
 }
 
+Box* setiter_next(Box* _self) noexcept {
+    RELEASE_ASSERT(_self->cls == set_iterator_cls, "");
+    BoxedSetIterator* self = (BoxedSetIterator*)_self;
+    if (!self->hasNext())
+        return NULL;
+    return self->next();
+}
+
 Box* setiteratorIter(BoxedSetIterator* self) {
     RELEASE_ASSERT(self->cls == set_iterator_cls, "");
     return self;
@@ -537,6 +545,7 @@ void setupSet() {
                                new BoxedFunction(boxRTFunction((void*)setiteratorHasnext, BOXED_BOOL, 1)));
     set_iterator_cls->giveAttr("next", new BoxedFunction(boxRTFunction((void*)setiteratorNext, UNKNOWN, 1)));
     set_iterator_cls->freeze();
+    set_iterator_cls->tp_iternext = setiter_next;
 
     set_cls->giveAttr("__new__",
                       new BoxedFunction(boxRTFunction((void*)setNew, UNKNOWN, 2, 1, false, false), { None }));

src/runtime/str.cpp

@@ -2401,6 +2401,17 @@ public:
         return characters[c & UCHAR_MAX];
     }
 
+    static Box* next_capi(Box* _self) noexcept {
+        assert(_self->cls == str_iterator_cls);
+        auto self = (BoxedStringIterator*)_self;
+        if (!hasnextUnboxed(self))
+            return NULL;
+
+        char c = *self->it;
+        ++self->it;
+        return characters[c & UCHAR_MAX];
+    }
+
     static void gcHandler(GCVisitor* v, Box* b) {
         Box::gcHandler(v, b);
         BoxedStringIterator* it = (BoxedStringIterator*)b;
@@ -2796,6 +2807,7 @@ void setupStr() {
     str_iterator_cls->giveAttr("next", new BoxedFunction(boxRTFunction((void*)BoxedStringIterator::next, STR, 1)));
     str_iterator_cls->freeze();
     str_iterator_cls->tpp_hasnext = (BoxedClass::pyston_inquiry)BoxedStringIterator::hasnextUnboxed;
+    str_iterator_cls->tp_iternext = BoxedStringIterator::next_capi;
 
     str_cls->tp_as_buffer = &string_as_buffer;
     str_cls->tp_print = string_print;

src/runtime/tuple.cpp

@@ -625,6 +625,7 @@ void setupTuple() {
 
     tuple_iterator_cls->freeze();
     tuple_iterator_cls->tpp_hasnext = tupleiterHasnextUnboxed;
+    tuple_iterator_cls->tp_iternext = tupleiter_next;
 }
 
 void teardownTuple() {

src/runtime/tuple.h

@@ -36,6 +36,7 @@ Box* tupleIter(Box* self) noexcept;
 Box* tupleIterIter(Box* self);
 Box* tupleiterHasnext(Box* self);
 i1 tupleiterHasnextUnboxed(Box* self);
+Box* tupleiter_next(Box* self) noexcept;
 Box* tupleiterNext(Box* self);
 }