Fix pyElements refcounting

src/core/types.h

@@ -562,18 +562,7 @@ public:
     BoxIteratorImpl* impl;
 
     BoxIterator(BoxIteratorImpl* impl) : impl(impl) {}
-    BoxIterator(const BoxIterator& rhs) = default;
-    BoxIterator(BoxIterator&& rhs) {
-        impl = rhs.impl;
-        rhs.impl = NULL;
-    }
-
-    ~BoxIterator() {
-        if (impl)
-            impl->~BoxIteratorImpl();
-    }
 
-    static llvm::iterator_range<BoxIterator> getRange(Box* container);
     bool operator==(BoxIterator const& rhs) const { return impl->isSame(rhs.impl); }
     bool operator!=(BoxIterator const& rhs) const { return !(*this == rhs); }
 
@@ -586,6 +575,23 @@ public:
     Box* operator*() { return impl->getValue(); }
 };
 
+// A custom "range" container that helps manage lifetimes.  We need to free the underlying Impl object
+// when the range loop is done; previously we had the iterator itself handle this, but that started
+// to get complicated since they get copied around, and the management of the begin() and end() iterators
+// is slightly different.
+// So to simplify, have the range object take care of it.
+class BoxIteratorRange {
+private:
+    std::unique_ptr<BoxIteratorImpl> begin_impl;
+    BoxIteratorImpl* end_impl;
+
+public:
+    BoxIteratorRange(std::unique_ptr<BoxIteratorImpl> begin, BoxIteratorImpl* end)
+        : begin_impl(std::move(begin)), end_impl(end) {}
+    BoxIterator begin() { return BoxIterator(begin_impl.get()); }
+    BoxIterator end() { return BoxIterator(end_impl); }
+};
+
 class HiddenClass;
 extern HiddenClass* root_hcls;
 
@@ -676,7 +682,7 @@ public:
     // Note: cls gets initialized in the new() function.
     BoxedClass* cls;
 
-    llvm::iterator_range<BoxIterator> pyElements();
+    BoxIteratorRange pyElements();
 
     // For instances with hc attrs:
     size_t getHCAttrsOffset();

src/runtime/builtin_modules/builtins.cpp

@@ -1195,13 +1195,13 @@ Box* zip(BoxedTuple* containers) {
     if (containers->size() == 0)
         return rtn;
 
-    std::vector<llvm::iterator_range<BoxIterator>> ranges;
+    std::vector<BoxIteratorRange> ranges;
     for (auto container : *containers) {
         ranges.push_back(container->pyElements());
     }
 
     std::vector<BoxIterator> iterators;
-    for (auto range : ranges) {
+    for (auto&& range : ranges) {
         iterators.push_back(range.begin());
     }
 
@@ -1307,13 +1307,14 @@ static BoxedClass* makeBuiltinException(BoxedClass* base, const char* name, int
 BoxedClass* enumerate_cls;
 class BoxedEnumerate : public Box {
 private:
+    BoxIteratorRange range;
     BoxIterator iterator, iterator_end;
     int64_t idx;
     BoxedLong* idx_long;
 
 public:
-    BoxedEnumerate(BoxIterator iterator_begin, BoxIterator iterator_end, int64_t idx, BoxedLong* idx_long)
-        : iterator(iterator_begin), iterator_end(iterator_end), idx(idx), idx_long(idx_long) {}
+    BoxedEnumerate(BoxIteratorRange range, int64_t idx, BoxedLong* idx_long)
+        : range(std::move(range)), iterator(range.begin()), iterator_end(range.end()), idx(idx), idx_long(idx_long) {}
 
     DEFAULT_CLASS(enumerate_cls);
 
@@ -1327,8 +1328,8 @@ public:
             assert(PyLong_Check(start));
             idx_long = (BoxedLong*)start;
         }
-        llvm::iterator_range<BoxIterator> range = obj->pyElements();
-        return new BoxedEnumerate(range.begin(), range.end(), idx, idx_long);
+        auto&& range = obj->pyElements();
+        return new BoxedEnumerate(std::move(range), idx, idx_long);
     }
 
     static Box* iter(Box* _self) noexcept {

src/runtime/iterators.cpp

@@ -99,8 +99,23 @@ public:
         }
     }
 
+    BoxIteratorIndex(const BoxIteratorIndex& rhs) : obj(rhs.obj), index(rhs.index) { Py_XINCREF(obj); }
+    BoxIteratorIndex(BoxIteratorIndex&& rhs) : obj(rhs.obj), index(rhs.index) { Py_XINCREF(obj); }
+    BoxIteratorIndex& operator=(const BoxIteratorIndex& rhs) {
+        obj = rhs.obj;
+        index = rhs.index;
+        Py_XINCREF(obj);
+        return *this;
+    }
+    BoxIteratorIndex& operator=(BoxIteratorIndex&& rhs) {
+        obj = rhs.obj;
+        index = rhs.index;
+        Py_XINCREF(obj);
+        return *this;
+    }
+
     ~BoxIteratorIndex() {
-        Py_XDECREF(obj);
+        Py_CLEAR(obj);
     }
 
     void next() override {
@@ -130,25 +145,26 @@ public:
 };
 }
 
-llvm::iterator_range<BoxIterator> BoxIterator::getRange(Box* container) {
-    if (container->cls == list_cls) {
+BoxIteratorRange Box::pyElements() {
+    if (this->cls == list_cls) {
         using BoxIteratorList = BoxIteratorIndex<BoxedList>;
-        BoxIterator begin = new BoxIteratorList((BoxedList*)container);
-        BoxIterator end = BoxIteratorList::end();
-        return llvm::iterator_range<BoxIterator>(std::move(begin), end);
-    } else if (container->cls == tuple_cls) {
+        std::unique_ptr<BoxIteratorImpl> begin(new BoxIteratorList((BoxedList*)this));
+        BoxIteratorImpl* end = BoxIteratorList::end();
+        return BoxIteratorRange(std::move(begin), end);
+    } else if (this->cls == tuple_cls) {
         using BoxIteratorTuple = BoxIteratorIndex<BoxedTuple>;
-        BoxIterator begin = new BoxIteratorTuple((BoxedTuple*)container);
-        BoxIterator end = BoxIteratorTuple::end();
-        return llvm::iterator_range<BoxIterator>(std::move(begin), end);
-    } else if (container->cls == str_cls) {
+        std::unique_ptr<BoxIteratorImpl> begin(new BoxIteratorTuple((BoxedTuple*)this));
+        BoxIteratorImpl* end = BoxIteratorTuple::end();
+        return BoxIteratorRange(std::move(begin), end);
+    } else if (this->cls == str_cls) {
         using BoxIteratorString = BoxIteratorIndex<BoxedString>;
-        BoxIterator begin = new BoxIteratorString((BoxedString*)container);
-        BoxIterator end = BoxIteratorString::end();
-        return llvm::iterator_range<BoxIterator>(std::move(begin), end);
+        std::unique_ptr<BoxIteratorImpl> begin(new BoxIteratorString((BoxedString*)this));
+        BoxIteratorImpl* end = BoxIteratorString::end();
+        return BoxIteratorRange(std::move(begin), end);
+    } else {
+        std::unique_ptr<BoxIteratorImpl> begin(new BoxIteratorGeneric(this));
+        BoxIteratorImpl* end = BoxIteratorGeneric::end();
+        return BoxIteratorRange(std::move(begin), end);
     }
-    BoxIterator begin = new BoxIteratorGeneric(container);
-    BoxIterator end = BoxIteratorGeneric::end();
-    return llvm::iterator_range<BoxIterator>(std::move(begin), end);
 }
 }

src/runtime/objmodel.cpp

@@ -6283,10 +6283,6 @@ Box* getiter(Box* o) {
     return getiterHelper(o);
 }
 
-llvm::iterator_range<BoxIterator> Box::pyElements() {
-    return BoxIterator::getRange(this);
-}
-
 void assertValidSlotIdentifier(Box* s) {
     // Ported from `valid_identifier` in cpython