Reduce code duplication with a template

src/runtime/iterators.cpp

@@ -12,33 +12,22 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "runtime/iterators.h"
-
+#include "core/types.h"
 #include "runtime/objmodel.h"
+#include "runtime/types.h"
 
 namespace pyston {
+namespace {
 
-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);
-    }
+static std::string next_str("next");
 
-    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);
-}
+class BoxIteratorGeneric : public BoxIteratorImpl {
+private:
+    Box* iterator;
+    Box* value;
 
-BoxIteratorGeneric::BoxIteratorGeneric(Box* container) : BoxIteratorImpl(container), iterator(nullptr), value(nullptr) {
+public:
+    BoxIteratorGeneric(Box* container) : BoxIteratorImpl(container), iterator(nullptr), value(nullptr) {
         if (container) {
             // TODO: this should probably call getPystonIter
             iterator = getiter(container);
@@ -47,10 +36,9 @@ BoxIteratorGeneric::BoxIteratorGeneric(Box* container) : BoxIteratorImpl(contain
             else
                 *this = end();
         }
-}
+    }
 
-static std::string next_str("next");
-void BoxIteratorGeneric::next() {
+    void next() override {
         assert(iterator);
         Box* hasnext = iterator->hasnextOrNullIC();
         if (hasnext) {
@@ -69,98 +57,83 @@ void BoxIteratorGeneric::next() {
                     throw e;
             }
         }
-}
+    }
 
-bool BoxIteratorGeneric::isSame(const BoxIteratorImpl* _rhs) {
-    BoxIteratorGeneric* rhs = (BoxIteratorGeneric*)_rhs;
-    return iterator == rhs->iterator && value == rhs->value;
-}
+    Box* getValue() override { return value; }
 
-void BoxIteratorGeneric::gcHandler(GCVisitor* v) {
+    void gcHandler(GCVisitor* v) override {
         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();
+    bool isSame(const BoxIteratorImpl* _rhs) override {
+        BoxIteratorGeneric* rhs = (BoxIteratorGeneric*)_rhs;
+        return iterator == rhs->iterator && value == rhs->value;
     }
-}
 
-Box* BoxIteratorList::getValue() {
-    return list->elts->elts[index];
-}
+    static BoxIteratorGeneric end() { return BoxIteratorGeneric(nullptr); }
+};
 
-void BoxIteratorList::gcHandler(GCVisitor* v) {
-    v->visitPotential(list);
-}
+template <typename T> class BoxIteratorIndex : public BoxIteratorImpl {
+private:
+    T* obj;
+    uint64_t index;
 
-bool BoxIteratorList::isSame(const BoxIteratorImpl* _rhs) {
-    BoxIteratorList* rhs = (BoxIteratorList*)_rhs;
-    return list == rhs->list && index == rhs->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])); }
 
-BoxIteratorTuple::BoxIteratorTuple(Box* container)
-    : BoxIteratorImpl(container), tuple((BoxedTuple*)container), index(0) {
-    assert(!container || container->cls == tuple_cls);
-    if (tuple && index >= tuple->elts.size())
+public:
+    BoxIteratorIndex(T* obj) : BoxIteratorImpl(obj), obj(obj), index(0) {
+        if (obj && !hasnext(obj, index))
             *this = end();
-}
+    }
 
-void BoxIteratorTuple::next() {
+    void next() override {
         if (!end().isSame(this)) {
             ++index;
-        if (index >= tuple->elts.size())
+            if (!hasnext(obj, index))
                 *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;
-}
+    Box* getValue() override { return getValue(obj, 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 gcHandler(GCVisitor* v) override { v->visitPotential(obj); }
 
-void BoxIteratorString::next() {
-    if (!end().isSame(this)) {
-        ++index;
-        if (index >= string->s.size())
-            *this = end();
+    bool isSame(const BoxIteratorImpl* _rhs) override {
+        auto rhs = (BoxIteratorIndex*)_rhs;
+        return obj == rhs->obj && index == rhs->index;
     }
-}
-
-Box* BoxIteratorString::getValue() {
-    return new BoxedString(std::string(1, string->s[index]));
-}
 
-void BoxIteratorString::gcHandler(GCVisitor* v) {
-    v->visitPotential(string);
+    static BoxIteratorIndex end() { return BoxIteratorIndex(nullptr); }
+};
 }
 
-bool BoxIteratorString::isSame(const BoxIteratorImpl* _rhs) {
-    BoxIteratorString* rhs = (BoxIteratorString*)_rhs;
-    return string == rhs->string && index == rhs->index;
+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/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