Add rewriting for 'a in b' expressions

src/codegen/unwinding.cpp

@@ -1007,7 +1007,7 @@ std::string getCurrentPythonLine() {
     return "unknown:-1";
 }
 
-void logByCurrentPythonLine(std::string& stat_name) {
+void logByCurrentPythonLine(const std::string& stat_name) {
     std::string stat = stat_name + "<" + getCurrentPythonLine() + ">";
     Stats::log(Stats::getStatCounter(stat));
 }

src/codegen/unwinding.h

@@ -46,7 +46,7 @@ ExecutionPoint getExecutionPoint();
 std::string getCurrentPythonLine();
 
 // doesn't really belong in unwinding.h, since it's stats related, but it needs to unwind to get the current line...
-void logByCurrentPythonLine(std::string& stat_name);
+void logByCurrentPythonLine(const std::string& stat_name);
 
 // Adds stack locals and closure locals into the locals dict, and returns it.
 Box* fastLocalsToBoxedLocals();

src/runtime/objmodel.cpp

@@ -3727,6 +3727,19 @@ static bool convert3wayCompareResultToBool(Box* v, int op_type) {
     };
 }
 
+Box* nonzeroAndBox(Box* b, bool negate) {
+    if (likely(b->cls == bool_cls)) {
+        if (negate)
+            return boxBool(b != True);
+        return b;
+    }
+
+    bool t = b->nonzeroIC();
+    if (negate)
+        t = !t;
+    return boxBool(t);
+}
+
 Box* compareInternal(Box* lhs, Box* rhs, int op_type, CompareRewriteArgs* rewrite_args) {
     if (op_type == AST_TYPE::Is || op_type == AST_TYPE::IsNot) {
         bool neg = (op_type == AST_TYPE::IsNot);
@@ -3742,11 +3755,26 @@ Box* compareInternal(Box* lhs, Box* rhs, int op_type, CompareRewriteArgs* rewrit
     }
 
     if (op_type == AST_TYPE::In || op_type == AST_TYPE::NotIn) {
-        // TODO do rewrite
-
         static BoxedString* contains_str = static_cast<BoxedString*>(PyString_InternFromString("__contains__"));
-        Box* contained = callattrInternal1(rhs, contains_str, CLASS_ONLY, NULL, ArgPassSpec(1), lhs);
+
+        Box* contained;
+        RewriterVar* r_contained;
+        if (rewrite_args) {
+            CallRewriteArgs crewrite_args(rewrite_args->rewriter, rewrite_args->rhs, rewrite_args->destination);
+            crewrite_args.arg1 = rewrite_args->lhs;
+            contained = callattrInternal1(rhs, contains_str, CLASS_ONLY, &crewrite_args, ArgPassSpec(1), lhs);
+
+            if (!crewrite_args.out_success)
+                rewrite_args = NULL;
+            else if (contained)
+                r_contained = crewrite_args.out_rtn;
+        } else {
+            contained = callattrInternal1(rhs, contains_str, CLASS_ONLY, NULL, ArgPassSpec(1), lhs);
+        }
+
         if (contained == NULL) {
+            rewrite_args = NULL;
+
             int result = _PySequence_IterSearch(rhs, lhs, PY_ITERSEARCH_CONTAINS);
             if (result < 0)
                 throwCAPIException();
@@ -3754,6 +3782,14 @@ Box* compareInternal(Box* lhs, Box* rhs, int op_type, CompareRewriteArgs* rewrit
             return boxBool(result);
         }
 
+        if (rewrite_args) {
+            auto r_negate = rewrite_args->rewriter->loadConst((int)(op_type == AST_TYPE::NotIn));
+            RewriterVar* r_contained_box
+                = rewrite_args->rewriter->call(true, (void*)nonzeroAndBox, r_contained, r_negate);
+            rewrite_args->out_rtn = r_contained_box;
+            rewrite_args->out_success = true;
+        }
+
         bool b;
         if (contained->cls == bool_cls)
             b = contained == True;

src/runtime/set.cpp

@@ -410,7 +410,7 @@ Box* setPop(BoxedSet* self) {
 
 Box* setContains(BoxedSet* self, Box* v) {
     RELEASE_ASSERT(PyAnySet_Check(self), "");
-    return boxBool(self->s.count(v) != 0);
+    return boxBool(self->s.find(v) != self->s.end());
 }
 
 Box* setEq(BoxedSet* self, BoxedSet* rhs) {