pyElements: Use specialiced iterators for known types

pyston (calibration)                      :    0.8s stock1: 0.8 (+0.8%)
pyston interp2.py                         :    5.9s stock1: 6.1 (-3.1%)
pyston raytrace.py                        :    7.5s stock1: 7.8 (-4.5%)
pyston nbody.py                           :   10.1s stock1: 9.9 (+1.5%)
pyston fannkuch.py                        :    6.7s stock1: 6.8 (-1.5%)
pyston chaos.py                           :   21.3s stock1: 22.6 (-5.8%)
pyston spectral_norm.py                   :   26.1s stock1: 27.9 (-6.7%)
pyston fasta.py                           :   17.7s stock1: 18.8 (-5.6%)
pyston pidigits.py                        :   16.1s stock1: 16.1 (+0.0%)
pyston richards.py                        :   10.2s stock1: 10.0 (+2.3%)
pyston deltablue.py                       :    2.4s stock1: 2.3 (+4.8%)
pyston (geomean-0b9f)                     :   10.1s stock1: 10.3 (-1.9%)

src/core/types.h

@@ -347,23 +347,37 @@ class NonzeroIC;
 class BinopIC;
 
 class Box;
-class BoxIterator {
-private:
-    Box* iter;
-    Box* value;
 
+class BoxIteratorImpl {
+public:
+    BoxIteratorImpl(Box* container) {}
+    virtual ~BoxIteratorImpl() = default;
+    virtual void next() = 0;
+    virtual Box* getValue() = 0;
+    virtual void gcHandler(GCVisitor* v) {}
+    virtual bool isSame(const BoxIteratorImpl* rhs) = 0;
+};
+
+class BoxIterator {
 public:
-    BoxIterator(Box* iter) : iter(iter), value(nullptr) {}
+    std::shared_ptr<BoxIteratorImpl> impl;
 
-    bool operator==(BoxIterator const& rhs) const { return (iter == rhs.iter && value == rhs.value); }
+    BoxIterator(std::shared_ptr<BoxIteratorImpl> impl) : impl(impl) {}
+    ~BoxIterator() = default;
+
+    static llvm::iterator_range<BoxIterator> getRange(Box* container);
+    bool operator==(BoxIterator const& rhs) const { return impl->isSame(rhs.impl.get()); }
     bool operator!=(BoxIterator const& rhs) const { return !(*this == rhs); }
 
-    BoxIterator& operator++();
+    BoxIterator& operator++() {
+        impl->next();
+        return *this;
+    }
 
-    Box* operator*() const { return value; }
-    Box* operator*() { return value; }
+    Box* operator*() const { return impl->getValue(); }
+    Box* operator*() { return impl->getValue(); }
 
-    void gcHandler(GCVisitor* v);
+    void gcHandler(GCVisitor* v) { impl->gcHandler(v); }
 };
 
 namespace gc {

src/runtime/iterators.cpp

+// Copyright (c) 2014-2015 Dropbox, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "runtime/iterators.h"
+
+#include "runtime/objmodel.h"
+
+namespace pyston {
+
+llvm::iterator_range<BoxIterator> BoxIterator::getRange(Box* container) {
+    if (container->cls == list_cls) {
+        BoxIterator begin(std::make_shared<BoxIteratorList>(container));
+        static BoxIterator end(std::make_shared<BoxIteratorList>(BoxIteratorList::end()));
+        return llvm::iterator_range<BoxIterator>(std::move(begin), end);
+    } else if (container->cls == tuple_cls) {
+        BoxIterator begin(std::make_shared<BoxIteratorTuple>(container));
+        static BoxIterator end(std::make_shared<BoxIteratorTuple>(BoxIteratorTuple::end()));
+        return llvm::iterator_range<BoxIterator>(std::move(begin), end);
+    } else if (container->cls == str_cls) {
+        BoxIterator begin(std::make_shared<BoxIteratorString>(container));
+        static BoxIterator end(std::make_shared<BoxIteratorString>(BoxIteratorString::end()));
+        return llvm::iterator_range<BoxIterator>(std::move(begin), end);
+    }
+
+    BoxIterator begin(std::make_shared<BoxIteratorGeneric>(container));
+    static BoxIterator end(std::make_shared<BoxIteratorGeneric>(BoxIteratorGeneric::end()));
+    return llvm::iterator_range<BoxIterator>(std::move(begin), end);
+}
+
+BoxIteratorGeneric::BoxIteratorGeneric(Box* container) : BoxIteratorImpl(container), iterator(nullptr), value(nullptr) {
+    if (container) {
+        // TODO: this should probably call getPystonIter
+        iterator = getiter(container);
+        if (iterator)
+            next();
+        else
+            *this = end();
+    }
+}
+
+static std::string next_str("next");
+void BoxIteratorGeneric::next() {
+    assert(iterator);
+    Box* hasnext = iterator->hasnextOrNullIC();
+    if (hasnext) {
+        if (hasnext->nonzeroIC()) {
+            value = iterator->nextIC();
+        } else {
+            *this = end();
+        }
+    } else {
+        try {
+            value = iterator->nextIC();
+        } catch (ExcInfo e) {
+            if (e.matches(StopIteration))
+                *this = end();
+            else
+                throw e;
+        }
+    }
+}
+
+bool BoxIteratorGeneric::isSame(const BoxIteratorImpl* _rhs) {
+    BoxIteratorGeneric* rhs = (BoxIteratorGeneric*)_rhs;
+    return iterator == rhs->iterator && value == rhs->value;
+}
+
+void BoxIteratorGeneric::gcHandler(GCVisitor* v) {
+    v->visitPotential(iterator);
+    v->visitPotential(value);
+}
+
+BoxIteratorList::BoxIteratorList(Box* container) : BoxIteratorImpl(container), list((BoxedList*)container), index(0) {
+    assert(!container || container->cls == list_cls);
+    if (list && index >= list->size)
+        *this = end();
+}
+
+void BoxIteratorList::next() {
+    if (!end().isSame(this)) {
+        ++index;
+        if (index >= list->size)
+            *this = end();
+    }
+}
+
+Box* BoxIteratorList::getValue() {
+    return list->elts->elts[index];
+}
+
+void BoxIteratorList::gcHandler(GCVisitor* v) {
+    v->visitPotential(list);
+}
+
+bool BoxIteratorList::isSame(const BoxIteratorImpl* _rhs) {
+    BoxIteratorList* rhs = (BoxIteratorList*)_rhs;
+    return list == rhs->list && index == rhs->index;
+}
+
+BoxIteratorTuple::BoxIteratorTuple(Box* container)
+    : BoxIteratorImpl(container), tuple((BoxedTuple*)container), index(0) {
+    assert(!container || container->cls == tuple_cls);
+    if (tuple && index >= tuple->elts.size())
+        *this = end();
+}
+
+void BoxIteratorTuple::next() {
+    if (!end().isSame(this)) {
+        ++index;
+        if (index >= tuple->elts.size())
+            *this = end();
+    }
+}
+
+Box* BoxIteratorTuple::getValue() {
+    return tuple->elts[index];
+}
+
+void BoxIteratorTuple::gcHandler(GCVisitor* v) {
+    v->visitPotential(tuple);
+}
+
+bool BoxIteratorTuple::isSame(const BoxIteratorImpl* _rhs) {
+    BoxIteratorTuple* rhs = (BoxIteratorTuple*)_rhs;
+    return tuple == rhs->tuple && index == rhs->index;
+}
+
+BoxIteratorString::BoxIteratorString(Box* container)
+    : BoxIteratorImpl(container), string((BoxedString*)container), index(0) {
+    assert(!container || container->cls == str_cls);
+    if (string && index >= string->s.size())
+        *this = end();
+}
+
+void BoxIteratorString::next() {
+    if (!end().isSame(this)) {
+        ++index;
+        if (index >= string->s.size())
+            *this = end();
+    }
+}
+
+Box* BoxIteratorString::getValue() {
+    return new BoxedString(std::string(1, string->s[index]));
+}
+
+void BoxIteratorString::gcHandler(GCVisitor* v) {
+    v->visitPotential(string);
+}
+
+bool BoxIteratorString::isSame(const BoxIteratorImpl* _rhs) {
+    BoxIteratorString* rhs = (BoxIteratorString*)_rhs;
+    return string == rhs->string && index == rhs->index;
+}
+}

src/runtime/iterators.h

+// Copyright (c) 2014-2015 Dropbox, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef PYSTON_RUNTIME_ITERATORS_H
+#define PYSTON_RUNTIME_ITERATORS_H
+
+#include "core/types.h"
+#include "runtime/types.h"
+
+namespace pyston {
+
+class BoxIteratorGeneric : public BoxIteratorImpl {
+private:
+    Box* iterator;
+    Box* value;
+
+public:
+    BoxIteratorGeneric(Box* container);
+
+    void next() override;
+    Box* getValue() override { return value; }
+    void gcHandler(GCVisitor* v) override;
+    bool isSame(const BoxIteratorImpl* rhs) override;
+
+    static BoxIteratorGeneric end() { return BoxIteratorGeneric(nullptr); }
+};
+
+class BoxIteratorList : public BoxIteratorImpl {
+private:
+    BoxedList* list;
+    int64_t index;
+
+public:
+    BoxIteratorList(Box* list);
+
+    void next() override;
+    Box* getValue() override;
+    void gcHandler(GCVisitor* v) override;
+    bool isSame(const BoxIteratorImpl* rhs) override;
+
+    static BoxIteratorList end() { return BoxIteratorList(nullptr); }
+};
+
+class BoxIteratorTuple : public BoxIteratorImpl {
+private:
+    BoxedTuple* tuple;
+    int64_t index;
+
+public:
+    BoxIteratorTuple(Box* tuple);
+
+    void next() override;
+    Box* getValue() override;
+    void gcHandler(GCVisitor* v) override;
+    bool isSame(const BoxIteratorImpl* rhs) override;
+
+    static BoxIteratorTuple end() { return BoxIteratorTuple(nullptr); }
+};
+
+class BoxIteratorString : public BoxIteratorImpl {
+private:
+    BoxedString* string;
+    int64_t index;
+
+public:
+    BoxIteratorString(Box* string);
+
+    void next() override;
+    Box* getValue() override;
+    void gcHandler(GCVisitor* v) override;
+    bool isSame(const BoxIteratorImpl* rhs) override;
+
+    static BoxIteratorString end() { return BoxIteratorString(nullptr); }
+};
+}
+
+#endif

src/runtime/objmodel.cpp

@@ -3463,10 +3463,7 @@ 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));
+    return BoxIterator::getRange(this);
 }
 
 // For use on __init__ return values

src/runtime/types.cpp

@@ -237,37 +237,6 @@ Box* Box::nextIC() {
     return this->cls->callNextIC(this);
 }
 
-BoxIterator& BoxIterator::operator++() {
-    static std::string next_str("next");
-
-    assert(iter);
-
-    Box* hasnext = iter->hasnextOrNullIC();
-    if (hasnext) {
-        if (hasnext->nonzeroIC()) {
-            value = iter->nextIC();
-        } else {
-            iter = nullptr;
-            value = nullptr;
-        }
-    } else {
-        try {
-            value = iter->nextIC();
-        } catch (ExcInfo e) {
-            if (e.matches(StopIteration)) {
-                iter = nullptr;
-                value = nullptr;
-            } else
-                throw e;
-        }
-    }
-    return *this;
-}
-
-void BoxIterator::gcHandler(GCVisitor* v) {
-    v->visitPotential(iter);
-    v->visitPotential(value);
-}
 
 std::string builtinStr("__builtin__");