Merge pull request #301 from undingen/perf_pyelements

pyElements: Use specialiced iterators for known types

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 "core/types.h"
+#include "runtime/objmodel.h"
+#include "runtime/types.h"
+
+namespace pyston {
+namespace {
+
+static std::string next_str("next");
+
+class BoxIteratorGeneric : public BoxIteratorImpl {
+private:
+    Box* iterator;
+    Box* value;
+
+public:
+    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();
+        }
+    }
+
+    void next() override {
+        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;
+            }
+        }
+    }
+
+    Box* getValue() override { return value; }
+
+    void gcHandler(GCVisitor* v) override {
+        v->visitPotential(iterator);
+        v->visitPotential(value);
+    }
+
+    bool isSame(const BoxIteratorImpl* _rhs) override {
+        BoxIteratorGeneric* rhs = (BoxIteratorGeneric*)_rhs;
+        return iterator == rhs->iterator && value == rhs->value;
+    }
+
+    static BoxIteratorGeneric end() { return BoxIteratorGeneric(nullptr); }
+};
+
+template <typename T> class BoxIteratorIndex : public BoxIteratorImpl {
+private:
+    T* obj;
+    uint64_t index;
+
+    static bool hasnext(BoxedList* o, uint64_t i) { return i < o->size; }
+    static Box* getValue(BoxedList* o, uint64_t i) { return o->elts->elts[i]; }
+
+    static bool hasnext(BoxedTuple* o, uint64_t i) { return i < o->elts.size(); }
+    static Box* getValue(BoxedTuple* o, uint64_t i) { return o->elts[i]; }
+
+    static bool hasnext(BoxedString* o, uint64_t i) { return i < o->s.size(); }
+    static Box* getValue(BoxedString* o, uint64_t i) { return new BoxedString(std::string(1, o->s[i])); }
+
+public:
+    BoxIteratorIndex(T* obj) : BoxIteratorImpl(obj), obj(obj), index(0) {
+        if (obj && !hasnext(obj, index))
+            *this = end();
+    }
+
+    void next() override {
+        if (!end().isSame(this)) {
+            ++index;
+            if (!hasnext(obj, index))
+                *this = end();
+        }
+    }
+
+    Box* getValue() override { return getValue(obj, index); }
+
+    void gcHandler(GCVisitor* v) override { v->visitPotential(obj); }
+
+    bool isSame(const BoxIteratorImpl* _rhs) override {
+        auto rhs = (BoxIteratorIndex*)_rhs;
+        return obj == rhs->obj && index == rhs->index;
+    }
+
+    static BoxIteratorIndex end() { return BoxIteratorIndex(nullptr); }
+};
+}
+
+llvm::iterator_range<BoxIterator> BoxIterator::getRange(Box* container) {
+    if (container->cls == list_cls) {
+        using BoxIteratorList = BoxIteratorIndex<BoxedList>;
+        BoxIterator begin(std::make_shared<BoxIteratorIndex<BoxedList>>((BoxedList*)container));
+        static BoxIterator end(std::make_shared<BoxIteratorIndex<BoxedList>>(BoxIteratorList::end()));
+        return llvm::iterator_range<BoxIterator>(std::move(begin), end);
+    } else if (container->cls == tuple_cls) {
+        using BoxIteratorTuple = BoxIteratorIndex<BoxedTuple>;
+        BoxIterator begin(std::make_shared<BoxIteratorIndex<BoxedTuple>>((BoxedTuple*)container));
+        static BoxIterator end(std::make_shared<BoxIteratorIndex<BoxedTuple>>(BoxIteratorTuple::end()));
+        return llvm::iterator_range<BoxIterator>(std::move(begin), end);
+    } else if (container->cls == str_cls) {
+        using BoxIteratorString = BoxIteratorIndex<BoxedString>;
+        BoxIterator begin(std::make_shared<BoxIteratorIndex<BoxedString>>((BoxedString*)container));
+        static BoxIterator end(std::make_shared<BoxIteratorIndex<BoxedString>>(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);
+}
+}

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__");