Implement chained comparisons

src/analysis/type_analysis.cpp

@@ -303,27 +303,29 @@ private:
     }
 
     virtual void* visit_compare(AST_Compare* node) {
-        RELEASE_ASSERT(node->ops.size() == 1, "unimplemented");
+        if (node->ops.size() == 1) {
+            CompilerType* left = getType(node->left);
+            CompilerType* right = getType(node->comparators[0]);
+
+            AST_TYPE::AST_TYPE op_type = node->ops[0];
+            if (op_type == AST_TYPE::Is || op_type == AST_TYPE::IsNot || op_type == AST_TYPE::In
+                || op_type == AST_TYPE::NotIn) {
+                assert(node->ops.size() == 1 && "I don't think this should happen");
+                return BOOL;
+            }
 
-        CompilerType* left = getType(node->left);
-        CompilerType* right = getType(node->comparators[0]);
+            const std::string& name = getOpName(node->ops[0]);
+            CompilerType* attr_type = left->getattrType(&name, true);
 
-        AST_TYPE::AST_TYPE op_type = node->ops[0];
-        if (op_type == AST_TYPE::Is || op_type == AST_TYPE::IsNot || op_type == AST_TYPE::In
-            || op_type == AST_TYPE::NotIn) {
-            assert(node->ops.size() == 1 && "I don't think this should happen");
-            return BOOL;
-        }
+            if (attr_type == UNDEF)
+                attr_type = UNKNOWN;
 
-        const std::string& name = getOpName(node->ops[0]);
-        CompilerType* attr_type = left->getattrType(&name, true);
-
-        if (attr_type == UNDEF)
-            attr_type = UNKNOWN;
-
-        std::vector<CompilerType*> arg_types;
-        arg_types.push_back(right);
-        return attr_type->callType(ArgPassSpec(2), arg_types, NULL);
+            std::vector<CompilerType*> arg_types;
+            arg_types.push_back(right);
+            return attr_type->callType(ArgPassSpec(2), arg_types, NULL);
+        } else {
+            return UNKNOWN;
+        }
     }
 
     virtual void* visit_dict(AST_Dict* node) {

src/core/cfg.cpp

@@ -486,18 +486,71 @@ private:
     }
 
     AST_expr* remapCompare(AST_Compare* node) {
-        AST_Compare* rtn = new AST_Compare();
-        rtn->lineno = node->lineno;
-        rtn->col_offset = node->col_offset;
+        // special case unchained comparisons to avoid generating a unnecessary complex cfg.
+        if (node->ops.size() == 1) {
+            AST_Compare* rtn = new AST_Compare();
+            rtn->lineno = node->lineno;
+            rtn->col_offset = node->col_offset;
 
-        rtn->ops = node->ops;
+            rtn->ops = node->ops;
 
-        rtn->left = remapExpr(node->left);
-        for (auto elt : node->comparators) {
-            rtn->comparators.push_back(remapExpr(elt));
-        }
+            rtn->left = remapExpr(node->left);
+            for (auto elt : node->comparators) {
+                rtn->comparators.push_back(remapExpr(elt));
+            }
+            return rtn;
+        } else {
+            std::string name = nodeName(node);
 
-        return rtn;
+            CFGBlock* exit_block = cfg->addDeferredBlock();
+            AST_expr* left = remapExpr(node->left);
+
+            for (int i = 0; i < node->ops.size(); i++) {
+                AST_expr* right = remapExpr(node->comparators[i]);
+
+                AST_Compare* val = new AST_Compare;
+                val->col_offset = node->col_offset;
+                val->lineno = node->lineno;
+                val->left = left;
+                val->comparators.push_back(right);
+                val->ops.push_back(node->ops[i]);
+
+                push_back(makeAssign(name, val));
+
+                AST_Branch* br = new AST_Branch();
+                br->test = val;
+                push_back(br);
+
+                CFGBlock* was_block = curblock;
+                CFGBlock* next_block = cfg->addBlock();
+                CFGBlock* crit_break_block = cfg->addBlock();
+                was_block->connectTo(next_block);
+                was_block->connectTo(crit_break_block);
+
+                br->iffalse = crit_break_block;
+                br->iftrue = next_block;
+
+                curblock = crit_break_block;
+                AST_Jump* j = new AST_Jump();
+                j->target = exit_block;
+                push_back(j);
+                crit_break_block->connectTo(exit_block);
+
+                curblock = next_block;
+
+                left = right;
+            }
+
+            AST_Jump* j = new AST_Jump();
+            push_back(j);
+            j->target = exit_block;
+            curblock->connectTo(exit_block);
+
+            cfg->placeBlock(exit_block);
+            curblock = exit_block;
+
+            return makeName(name, AST_TYPE::Load);
+        }
     }
 
     AST_expr* remapDict(AST_Dict* node) {

test/tests/comparisons_chained.py

-# expected: fail
-# - chained comparisons
-# - "in" and "not in"
-
 def f(n):
     print "f(%d)" % n
     return n