Call sq_contains for extension types

src/capi/typeobject.cpp

@@ -1078,7 +1078,7 @@ static int slot_sq_ass_slice(PyObject* self, Py_ssize_t i, Py_ssize_t j, PyObjec
     return 0;
 }
 
-static int slot_sq_contains(PyObject* self, PyObject* value) noexcept {
+/* Pyston change: static*/ int slot_sq_contains(PyObject* self, PyObject* value) noexcept {
     STAT_TIMER(t0, "us_timer_slot_sqcontains", SLOT_AVOIDABILITY(self));
 
     PyObject* func, *res, *args;

src/capi/typeobject.h

@@ -38,6 +38,7 @@ PyObject* slot_tp_richcompare(PyObject* self, PyObject* other, int op) 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;
+int slot_sq_contains(PyObject* self, PyObject* value) noexcept;
 }
 
 #endif

src/runtime/inline/boxing.h

@@ -37,6 +37,12 @@ extern "C" inline Box* boxBool(bool b) {
     return rtn;
 }
 
+extern "C" inline Box* boxBoolNegated(bool b) __attribute__((visibility("default")));
+extern "C" inline Box* boxBoolNegated(bool b) {
+    Box* rtn = b ? False : True;
+    return rtn;
+}
+
 extern "C" inline bool unboxBool(Box* b) __attribute__((visibility("default")));
 extern "C" inline bool unboxBool(Box* b) {
     assert(b->cls == bool_cls);

src/runtime/inline/link_forcer.cpp

@@ -57,6 +57,7 @@ void force() {
     FORCE(unboxCLFunction);
     FORCE(boxInstanceMethod);
     FORCE(boxBool);
+    FORCE(boxBoolNegated);
     FORCE(unboxBool);
     FORCE(createTuple);
     FORCE(createDict);

src/runtime/objmodel.cpp

@@ -4008,6 +4008,46 @@ Box* compareInternal(Box* lhs, Box* rhs, int op_type, CompareRewriteArgs* rewrit
     if (op_type == AST_TYPE::In || op_type == AST_TYPE::NotIn) {
         static BoxedString* contains_str = static_cast<BoxedString*>(PyString_InternFromString("__contains__"));
 
+        // The checks for this branch are taken from CPython's PySequence_Contains
+        if (PyType_HasFeature(rhs->cls, Py_TPFLAGS_HAVE_SEQUENCE_IN)) {
+            PySequenceMethods* sqm = rhs->cls->tp_as_sequence;
+            if (sqm != NULL && sqm->sq_contains != NULL && sqm->sq_contains != slot_sq_contains) {
+                if (rewrite_args) {
+                    RewriterVar* r_lhs = rewrite_args->lhs;
+                    RewriterVar* r_rhs = rewrite_args->rhs;
+                    RewriterVar* r_cls = r_rhs->getAttr(offsetof(Box, cls));
+                    RewriterVar* r_sqm = r_cls->getAttr(offsetof(BoxedClass, tp_as_sequence));
+                    r_sqm->addGuardNotEq(0);
+                    // We might need to guard on tp_flags if they can change?
+
+                    // Currently, guard that the value of sq_contains didn't change, and then
+                    // emit a call to the current function address.
+                    // It might be better to just load the current value of sq_contains and call it
+                    // (after guarding it's not null), or maybe not.  But the rewriter doesn't currently
+                    // support calling a RewriterVar (can only call fixed function addresses).
+                    r_sqm->addAttrGuard(offsetof(PySequenceMethods, sq_contains), (intptr_t)sqm->sq_contains);
+                    RewriterVar* r_b = rewrite_args->rewriter->call(true, (void*)sqm->sq_contains, r_rhs, r_lhs);
+                    rewrite_args->rewriter->call(true, (void*)checkAndThrowCAPIException);
+                    // This could be inlined:
+                    RewriterVar* r_r;
+                    if (op_type == AST_TYPE::NotIn)
+                        r_r = rewrite_args->rewriter->call(false, (void*)boxBoolNegated, r_b);
+                    else
+                        r_r = rewrite_args->rewriter->call(false, (void*)boxBool, r_b);
+
+                    rewrite_args->out_success = true;
+                    rewrite_args->out_rtn = r_r;
+                }
+
+                int r = (*sqm->sq_contains)(rhs, lhs);
+                if (r == -1)
+                    throwCAPIException();
+                if (op_type == AST_TYPE::NotIn)
+                    r = !r;
+                return boxBool(r);
+            }
+        }
+
         Box* contained;
         RewriterVar* r_contained;
         if (rewrite_args) {

src/runtime/types.h

@@ -106,6 +106,7 @@ extern BoxedModule* sys_module, *builtins_module, *math_module, *time_module, *t
 }
 
 extern "C" Box* boxBool(bool);
+extern "C" Box* boxBoolNegated(bool);
 extern "C" Box* boxInt(i64) __attribute__((visibility("default")));
 extern "C" i64 unboxInt(Box*);
 extern "C" Box* boxFloat(double d);