Add refcounts to enumerate()

src/core/types.h

@@ -549,12 +549,21 @@ class BinopIC;
 
 class Box;
 
+#define Py_TRAVERSE(obj)                                                                                               \
+    do {                                                                                                               \
+        int vret = (obj).traverse(visit, arg);                                                                         \
+        if (vret)                                                                                                      \
+            return vret;                                                                                               \
+    } while (0)
+
+
 class BoxIteratorImpl {
 public:
     virtual ~BoxIteratorImpl() = default;
     virtual void next() = 0;
     virtual Box* getValue() = 0;
     virtual bool isSame(const BoxIteratorImpl* rhs) = 0;
+    virtual int traverse(visitproc visit, void* arg) = 0;
 };
 
 class BoxIterator {
@@ -590,6 +599,12 @@ public:
         : begin_impl(std::move(begin)), end_impl(end) {}
     BoxIterator begin() { return BoxIterator(begin_impl.get()); }
     BoxIterator end() { return BoxIterator(end_impl); }
+
+    int traverse(visitproc visit, void* arg) {
+        Py_TRAVERSE(*begin_impl);
+        Py_TRAVERSE(*end_impl);
+        return 0;
+    }
 };
 
 class HiddenClass;
@@ -637,14 +652,6 @@ public:
 };
 static_assert(sizeof(HCAttrs) == sizeof(struct _hcattrs), "");
 
-#define Py_VISIT_HCATTRS(hcattrs)                                                                                      \
-    do {                                                                                                               \
-        int vret = hcattrs.traverse(visit, arg);                                                                      \
-        if (vret)                                                                                                      \
-            return vret;                                                                                               \
-    } while (0)
-
-
 extern std::vector<BoxedClass*> classes;
 
 // Debugging helper: pass this as a tp_clear function to say that you have explicitly verified

src/runtime/builtin_modules/builtins.cpp

@@ -1330,7 +1330,13 @@ private:
 
 public:
     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) {}
+        : range(std::move(range)),
+          iterator(this->range.begin()),
+          iterator_end(this->range.end()),
+          idx(idx),
+          idx_long(idx_long) {
+        Py_XINCREF(idx_long);
+    }
 
     DEFAULT_CLASS(enumerate_cls);
 
@@ -1358,8 +1364,13 @@ public:
         assert(_self->cls == enumerate_cls);
         BoxedEnumerate* self = static_cast<BoxedEnumerate*>(_self);
         Box* val = *self->iterator;
+        AUTO_DECREF(val);
         ++self->iterator;
-        Box* rtn = BoxedTuple::create({ self->idx_long ? self->idx_long : boxInt(self->idx), val });
+        Box* rtn;
+        if (self->idx_long)
+            rtn = BoxedTuple::create({ self->idx_long, val });
+        else
+            rtn = BoxedTuple::create({ autoDecref(boxInt(self->idx)), val });
 
         // check if incrementing the counter would overflow it, if so switch to long counter
         if (self->idx == PY_SSIZE_T_MAX) {
@@ -1368,7 +1379,7 @@ public:
             self->idx = -1;
         }
         if (self->idx_long)
-            self->idx_long = (BoxedLong*)longAdd(self->idx_long, boxInt(1));
+            self->idx_long = (BoxedLong*)longAdd(autoDecref(self->idx_long), autoDecref(boxInt(1)));
         else
             ++self->idx;
         return rtn;
@@ -1381,10 +1392,27 @@ public:
     }
 
     static void dealloc(Box* b) noexcept {
-        Py_FatalError("unimplemented");
+        assert(b->cls == enumerate_cls);
+        BoxedEnumerate* self = static_cast<BoxedEnumerate*>(b);
+
+        PyObject_GC_UnTrack(self);
+
+        Py_XDECREF(self->idx_long);
+        self->iterator.~BoxIterator();
+        self->iterator_end.~BoxIterator();
+        self->range.~BoxIteratorRange();
+
+        self->cls->tp_free(self);
     }
-    static int traverse(Box* self, visitproc visit, void *arg) noexcept {
-        Py_FatalError("unimplemented");
+
+    static int traverse(Box* b, visitproc visit, void *arg) noexcept {
+        assert(b->cls == enumerate_cls);
+        BoxedEnumerate* self = static_cast<BoxedEnumerate*>(b);
+
+        Py_VISIT(self->idx_long);
+        Py_TRAVERSE(self->range);
+
+        return 0;
     }
 };
 

src/runtime/classobj.cpp

@@ -1754,7 +1754,7 @@ void BoxedInstance::dealloc(Box* b) noexcept {
 int BoxedInstance::traverse(Box* o, visitproc visit, void* arg) noexcept {
     BoxedInstance* self = static_cast<BoxedInstance*>(o);
 
-    Py_VISIT_HCATTRS(self->attrs);
+    Py_TRAVERSE(self->attrs);
     Py_VISIT(self->inst_cls);
 
     return 0;
@@ -1785,7 +1785,7 @@ int BoxedClassobj::traverse(Box* o, visitproc visit, void* arg) noexcept {
 
     Py_VISIT(cl->bases);
     Py_VISIT(cl->name);
-    Py_VISIT_HCATTRS(cl->attrs);
+    Py_TRAVERSE(cl->attrs);
     return 0;
 }
 

src/runtime/descr.cpp

@@ -774,7 +774,6 @@ template <ExceptionStyle S>
 Box* BoxedWrapperObject::tppCall(Box* _self, CallRewriteArgs* rewrite_args, ArgPassSpec argspec, Box* arg1, Box* arg2,
                                  Box* arg3, Box** args,
                                  const std::vector<BoxedString*>* keyword_names) noexcept(S == CAPI) {
-    assert(0 && "check refcounting");
     STAT_TIMER(t0, "us_timer_boxedwrapperobject_call", (_self->cls->is_user_defined ? 10 : 20));
 
     assert(_self->cls == wrapperobject_cls);

src/runtime/iterators.cpp

@@ -69,6 +69,13 @@ public:
         return iterator == rhs->iterator && value == rhs->value;
     }
 
+    int traverse(visitproc visit, void* arg) override {
+        Py_VISIT(iterator);
+        Py_VISIT(value);
+
+        return 0;
+    }
+
     static BoxIteratorGeneric* end() {
         static BoxIteratorGeneric _end(nullptr);
         return &_end;
@@ -138,6 +145,12 @@ public:
         return obj == rhs->obj && index == rhs->index;
     }
 
+    int traverse(visitproc visit, void* arg) override {
+        Py_VISIT(obj);
+
+        return 0;
+    }
+
     static BoxIteratorIndex* end() {
         static BoxIteratorIndex _end(nullptr);
         return &_end;

src/runtime/types.cpp

@@ -428,7 +428,7 @@ static int func_traverse(BoxedFunction* f, visitproc visit, void* arg) noexcept
     Py_VISIT(f->closure);
 
     //Py_VISIT(f->func_dict);
-    Py_VISIT_HCATTRS(f->attrs);
+    Py_TRAVERSE(f->attrs);
     return 0;
 }
 
@@ -3773,7 +3773,7 @@ void BoxedModule::dealloc(Box* b) noexcept {
 
 int BoxedModule::traverse(Box* _m, visitproc visit, void* arg) noexcept {
     BoxedModule* m = static_cast<BoxedModule*>(_m);
-    Py_VISIT_HCATTRS(m->attrs);
+    Py_TRAVERSE(m->attrs);
     assert(!m->keep_alive.size());
     return 0;
 }
@@ -3913,7 +3913,7 @@ type_traverse(PyTypeObject *type, visitproc visit, void *arg)
     // Pyston change: HEAPTYPE is not about whether it is in GC or not
     // assert(type->tp_flags & Py_TPFLAGS_HEAPTYPE);
 
-    Py_VISIT_HCATTRS(type->attrs);
+    Py_TRAVERSE(type->attrs);
 
     Py_VISIT(type->tp_dict);
     Py_VISIT(type->tp_cache);

test/tests/enumerate_test.py

+import sys
+print list(enumerate(range(100)))
+print list(enumerate(range(100), sys.maxint-50))
+
+# cycle collection:
+print enumerate(range(100)).next()