Merge pull request #794 from undingen/tp_iter

getiter: use tp_iter for non python types

src/capi/typeobject.cpp

@@ -807,7 +807,7 @@ static PyObject* half_richcompare(PyObject* self, PyObject* other, int op) noexc
     return Py_NotImplemented;
 }
 
-static PyObject* slot_tp_iter(PyObject* self) noexcept {
+PyObject* slot_tp_iter(PyObject* self) noexcept {
     STAT_TIMER(t0, "us_timer_slot_tpiter", SLOT_AVOIDABILITY(self));
 
     PyObject* func, *res;

src/capi/typeobject.h

@@ -43,6 +43,7 @@ PyObject* mro_external(PyObject* self) noexcept;
 int type_set_bases(PyTypeObject* type, PyObject* value, void* context) noexcept;
 
 PyObject* slot_tp_richcompare(PyObject* self, PyObject* other, int op) noexcept;
+PyObject* slot_tp_iter(PyObject* self) noexcept;
 PyObject* slot_tp_iternext(PyObject* self) noexcept;
 PyObject* slot_tp_new(PyTypeObject* self, PyObject* args, PyObject* kwds) noexcept;
 PyObject* slot_mp_subscript(PyObject* self, PyObject* arg1) noexcept;

src/runtime/builtin_modules/builtins.cpp

@@ -992,7 +992,7 @@ public:
         return new BoxedEnumerate(range.begin(), range.end(), idx);
     }
 
-    static Box* iter(Box* _self) {
+    static Box* iter(Box* _self) noexcept {
         assert(_self->cls == enumerate_cls);
         BoxedEnumerate* self = static_cast<BoxedEnumerate*>(_self);
         return self;
@@ -1478,6 +1478,7 @@ void setupBuiltins() {
     enumerate_cls->giveAttr("__hasnext__",
                             new BoxedFunction(boxRTFunction((void*)BoxedEnumerate::hasnext, BOXED_BOOL, 1)));
     enumerate_cls->freeze();
+    enumerate_cls->tp_iter = BoxedEnumerate::iter;
     builtins_module->giveAttr("enumerate", enumerate_cls);
 
 

src/runtime/dict.cpp

@@ -491,6 +491,12 @@ extern "C" int PyDict_Contains(PyObject* op, PyObject* key) noexcept {
 
     try {
         if (op->cls == attrwrapper_cls) {
+            if (key->cls == str_cls) {
+                BoxedString* key_str = (BoxedString*)key;
+                internStringMortalInplace(key_str);
+                return unwrapAttrWrapper(op)->hasattr(key_str);
+            }
+
             Box* rtn = PyObject_CallMethod(op, "__contains__", "O", key);
             if (!rtn)
                 return -1;

src/runtime/generator.cpp

@@ -455,5 +455,6 @@ void setupGenerator() {
     generator_cls->giveAttr("__name__", new (pyston_getset_cls) BoxedGetsetDescriptor(generatorName, NULL, NULL));
 
     generator_cls->freeze();
+    generator_cls->tp_iter = PyObject_SelfIter;
 }
 }

src/runtime/inline/list.cpp

@@ -29,7 +29,7 @@ Box* listIterIter(Box* s) {
     return s;
 }
 
-Box* listIter(Box* s) {
+Box* listIter(Box* s) noexcept {
     assert(isSubclass(s->cls, list_cls));
     BoxedList* self = static_cast<BoxedList*>(s);
     return new BoxedListIterator(self, 0);

src/runtime/inline/tuple.cpp

@@ -26,7 +26,7 @@ Box* tupleIterIter(Box* s) {
     return s;
 }
 
-Box* tupleIter(Box* s) {
+Box* tupleIter(Box* s) noexcept {
     assert(isSubclass(s->cls, tuple_cls));
     BoxedTuple* self = static_cast<BoxedTuple*>(s);
     return new BoxedTupleIterator(self);

src/runtime/inline/xrange.cpp

@@ -158,7 +158,7 @@ Box* xrangeIterIter(Box* self) {
     return self;
 }
 
-Box* xrangeIter(Box* self) {
+Box* xrangeIter(Box* self) noexcept {
     assert(self->cls == xrange_cls);
 
     Box* rtn = new BoxedXrangeIterator(static_cast<BoxedXrange*>(self), false);
@@ -224,6 +224,8 @@ void setupXrange() {
     xrange_cls->giveAttr("__iterator_cls__", xrange_iterator_cls);
 
     xrange_cls->freeze();
+    xrange_cls->tp_iter = xrangeIter;
+
     xrange_iterator_cls->freeze();
     xrange_iterator_cls->tpp_hasnext = BoxedXrangeIterator::xrangeIteratorHasnextUnboxed;
 }

src/runtime/list.cpp

@@ -1183,6 +1183,7 @@ void setupList() {
 
     list_cls->tp_as_sequence->sq_slice = list_slice;
     list_cls->tp_as_sequence->sq_length = list_length;
+    list_cls->tp_iter = listIter;
 
     CLFunction* hasnext = boxRTFunction((void*)listiterHasnextUnboxed, BOOL, 1);
     addRTFunction(hasnext, (void*)listiterHasnext, BOXED_BOOL);

src/runtime/list.h

@@ -31,7 +31,7 @@ public:
     DEFAULT_CLASS(list_iterator_cls);
 };
 
-Box* listIter(Box* self);
+Box* listIter(Box* self) noexcept;
 Box* listIterIter(Box* self);
 Box* listiterHasnext(Box* self);
 i1 listiterHasnextUnboxed(Box* self);

src/runtime/objmodel.cpp

@@ -4912,8 +4912,14 @@ extern "C" Box* getiterHelper(Box* o) {
 
 Box* getiter(Box* o) {
     // TODO add rewriting to this?  probably want to try to avoid this path though
-    static BoxedString* iter_str = internStringImmortal("__iter__");
-    Box* r = callattrInternal0(o, iter_str, LookupScope::CLASS_ONLY, NULL, ArgPassSpec(0));
+    BoxedClass* type = o->cls;
+    Box* r = NULL;
+    if (PyType_HasFeature(type, Py_TPFLAGS_HAVE_ITER) && type->tp_iter != slot_tp_iter && type->tp_iter) {
+        r = type->tp_iter(o);
+    } else {
+        static BoxedString* iter_str = internStringImmortal("__iter__");
+        r = callattrInternal0(o, iter_str, LookupScope::CLASS_ONLY, NULL, ArgPassSpec(0));
+    }
     if (r) {
         if (!PyIter_Check(r)) {
             raiseExcHelper(TypeError, "iter() returned non-iterator of type '%s'", r->cls->tp_name);

src/runtime/set.cpp

@@ -214,7 +214,7 @@ Box* setXorSet(BoxedSet* lhs, BoxedSet* rhs) {
     return rtn;
 }
 
-Box* setIter(BoxedSet* self) {
+Box* setIter(BoxedSet* self) noexcept {
     RELEASE_ASSERT(PyAnySet_Check(self), "");
     return new BoxedSetIterator(self);
 }
@@ -604,6 +604,9 @@ void setupSet() {
 
     set_cls->freeze();
     frozenset_cls->freeze();
+
+    set_cls->tp_iter = (decltype(set_cls->tp_iter))setIter;
+    frozenset_cls->tp_iter = (decltype(frozenset_cls->tp_iter))setIter;
 }
 
 void teardownSet() {

src/runtime/str.cpp

@@ -2311,7 +2311,7 @@ extern "C" void strIteratorGCHandler(GCVisitor* v, Box* b) {
     v->visit(it->s);
 }
 
-Box* strIter(BoxedString* self) {
+Box* strIter(BoxedString* self) noexcept {
     assert(PyString_Check(self));
     return new BoxedStringIterator(self);
 }
@@ -2791,6 +2791,7 @@ void setupStr() {
 
     str_cls->tp_as_sequence->sq_slice = str_slice;
     str_cls->tp_as_sequence->sq_length = str_length;
+    str_cls->tp_iter = (decltype(str_cls->tp_iter))strIter;
 
     basestring_cls->giveAttr("__doc__",
                              boxString("Type basestring cannot be instantiated; it is the base for str and unicode."));

src/runtime/tuple.cpp

@@ -605,6 +605,7 @@ void setupTuple() {
 
     tuple_cls->tp_as_sequence->sq_item = (ssizeargfunc)tupleitem;
     tuple_cls->tp_as_sequence->sq_length = (lenfunc)tuplelength;
+    tuple_cls->tp_iter = tupleIter;
 
     CLFunction* hasnext = boxRTFunction((void*)tupleiterHasnextUnboxed, BOOL, 1);
     addRTFunction(hasnext, (void*)tupleiterHasnext, BOXED_BOOL);

src/runtime/tuple.h

@@ -30,7 +30,7 @@ public:
     DEFAULT_CLASS(tuple_iterator_cls);
 };
 
-Box* tupleIter(Box* self);
+Box* tupleIter(Box* self) noexcept;
 Box* tupleIterIter(Box* self);
 Box* tupleiterHasnext(Box* self);
 i1 tupleiterHasnextUnboxed(Box* self);