Merge pull request #714 from kmod/perf4

Convert "a in (b, c)" to "a == b or a == c"

src/codegen/compvars.cpp

@@ -455,6 +455,10 @@ public:
         return new ConcreteCompilerVariable(UNKNOWN, rtn, true);
     }
 
+    CompilerVariable* contains(IREmitter& emitter, const OpInfo& info, VAR* var, CompilerVariable* lhs) override {
+        return lhs->binexp(emitter, info, var, AST_TYPE::In, Compare);
+    }
+
     Box* deserializeFromFrame(const FrameVals& vals) override {
         assert(vals.size() == 1);
         return reinterpret_cast<Box*>(vals[0]);
@@ -1081,6 +1085,10 @@ public:
         }
     }
 
+    CompilerVariable* contains(IREmitter& emitter, const OpInfo& info, VAR* var, CompilerVariable* lhs) override {
+        return makeBool(false);
+    }
+
     ConcreteCompilerType* getBoxType() override { return BOXED_INT; }
 
     Box* deserializeFromFrame(const FrameVals& vals) override {
@@ -1319,6 +1327,10 @@ public:
         return rtn;
     }
 
+    CompilerVariable* contains(IREmitter& emitter, const OpInfo& info, VAR* var, CompilerVariable* lhs) override {
+        return makeBool(false);
+    }
+
     ConcreteCompilerType* getBoxType() override { return BOXED_FLOAT; }
 
     Box* deserializeFromFrame(const FrameVals& vals) override {
@@ -1688,6 +1700,10 @@ public:
         return rtn;
     }
 
+    CompilerVariable* contains(IREmitter& emitter, const OpInfo& info, VAR* var, CompilerVariable* lhs) override {
+        return UNKNOWN->contains(emitter, info, var, lhs);
+    }
+
     CompilerVariable* getitem(IREmitter& emitter, const OpInfo& info, VAR* var, CompilerVariable* slice) override {
         static BoxedString* attr = static_cast<BoxedString*>(PyString_InternFromString("__getitem__"));
         bool no_attribute = false;
@@ -1914,6 +1930,13 @@ public:
         return rtn;
     }
 
+    CompilerVariable* contains(IREmitter& emitter, const OpInfo& info, VAR* var, CompilerVariable* lhs) override {
+        ConcreteCompilerVariable* converted = var->makeConverted(emitter, STR);
+        CompilerVariable* rtn = converted->contains(emitter, info, lhs);
+        converted->decvref(emitter);
+        return rtn;
+    }
+
     ConcreteCompilerVariable* nonzero(IREmitter& emitter, const OpInfo& info, VAR* var) override {
         return makeBool(var->getValue()->size() != 0);
     }
@@ -2033,6 +2056,10 @@ public:
         return rtn;
     }
 
+    CompilerVariable* contains(IREmitter& emitter, const OpInfo& info, VAR* var, CompilerVariable* lhs) override {
+        return makeBool(false);
+    }
+
     ConcreteCompilerType* getBoxType() override { return BOXED_BOOL; }
 
     Box* deserializeFromFrame(const FrameVals& vals) override {
@@ -2047,6 +2074,23 @@ ConcreteCompilerVariable* makeBool(bool b) {
     return new ConcreteCompilerVariable(BOOL, llvm::ConstantInt::get(BOOL->llvmType(), b, false), true);
 }
 
+ConcreteCompilerVariable* doIs(IREmitter& emitter, CompilerVariable* lhs, CompilerVariable* rhs, bool negate) {
+    // TODO: I think we can do better here and not force the types to box themselves
+
+    ConcreteCompilerVariable* converted_left = lhs->makeConverted(emitter, UNKNOWN);
+    ConcreteCompilerVariable* converted_right = rhs->makeConverted(emitter, UNKNOWN);
+    llvm::Value* cmp;
+    if (!negate)
+        cmp = emitter.getBuilder()->CreateICmpEQ(converted_left->getValue(), converted_right->getValue());
+    else
+        cmp = emitter.getBuilder()->CreateICmpNE(converted_left->getValue(), converted_right->getValue());
+
+    converted_left->decvref(emitter);
+    converted_right->decvref(emitter);
+
+    return boolFromI1(emitter, cmp);
+}
+
 ConcreteCompilerType* BOXED_TUPLE;
 class TupleType : public ValuedCompilerType<const std::vector<CompilerVariable*>*> {
 private:
@@ -2196,6 +2240,58 @@ public:
         return rtn;
     }
 
+    CompilerVariable* contains(IREmitter& emitter, const OpInfo& info, VAR* var, CompilerVariable* lhs) override {
+        llvm::SmallVector<std::pair<llvm::BasicBlock*, llvm::Value*>, 4> phi_incoming;
+
+        llvm::BasicBlock* end = emitter.createBasicBlock();
+
+        ConcreteCompilerVariable* converted_lhs = lhs->makeConverted(emitter, lhs->getConcreteType());
+
+        for (CompilerVariable* e : *var->getValue()) {
+            // TODO: we could potentially avoid the identity tests if we know that either type has
+            // an __eq__ that is reflexive (returns True for the same object).
+            {
+                ConcreteCompilerVariable* is_same = doIs(emitter, converted_lhs, e, false);
+                llvm::Value* raw = i1FromBool(emitter, is_same);
+
+                phi_incoming.push_back(std::make_pair(emitter.currentBasicBlock(), getConstantInt(1, g.i1)));
+                llvm::BasicBlock* new_bb = emitter.createBasicBlock();
+                new_bb->moveAfter(emitter.currentBasicBlock());
+                emitter.getBuilder()->CreateCondBr(raw, end, new_bb);
+                emitter.setCurrentBasicBlock(new_bb);
+            }
+
+            {
+                CompilerVariable* eq = converted_lhs->binexp(emitter, info, e, AST_TYPE::Eq, Compare);
+                ConcreteCompilerVariable* eq_nonzero = eq->nonzero(emitter, info);
+                assert(eq_nonzero->getType() == BOOL);
+                llvm::Value* raw = i1FromBool(emitter, eq_nonzero);
+
+                phi_incoming.push_back(std::make_pair(emitter.currentBasicBlock(), getConstantInt(1, g.i1)));
+                llvm::BasicBlock* new_bb = emitter.createBasicBlock();
+                new_bb->moveAfter(emitter.currentBasicBlock());
+                emitter.getBuilder()->CreateCondBr(raw, end, new_bb);
+                emitter.setCurrentBasicBlock(new_bb);
+            }
+        }
+
+        // TODO This last block is unnecessary:
+        phi_incoming.push_back(std::make_pair(emitter.currentBasicBlock(), getConstantInt(0, g.i1)));
+        emitter.getBuilder()->CreateBr(end);
+
+        end->moveAfter(emitter.currentBasicBlock());
+        emitter.setCurrentBasicBlock(end);
+
+        auto phi = emitter.getBuilder()->CreatePHI(g.i1, phi_incoming.size());
+        for (auto p : phi_incoming) {
+            phi->addIncoming(p.second, p.first);
+        }
+
+        converted_lhs->decvref(emitter);
+
+        return boolFromI1(emitter, phi);
+    }
+
     CompilerVariable* callattr(IREmitter& emitter, const OpInfo& info, VAR* var, BoxedString* attr, CallattrFlags flags,
                                const std::vector<CompilerVariable*>& args,
                                const std::vector<BoxedString*>* keyword_names) override {
@@ -2320,6 +2416,10 @@ public:
         return undefVariable();
     }
 
+    CompilerVariable* contains(IREmitter& emitter, const OpInfo& info, VAR* var, CompilerVariable* lhs) override {
+        return undefVariable();
+    }
+
     CompilerVariable* getitem(IREmitter& emitter, const OpInfo& info, ConcreteCompilerVariable* var,
                               CompilerVariable* slice) override {
         return undefVariable();

src/codegen/compvars.h

@@ -152,6 +152,7 @@ public:
         printf("binexp not defined for %s\n", debugName().c_str());
         abort();
     }
+    virtual CompilerVariable* contains(IREmitter& emitter, const OpInfo& info, VAR* var, CompilerVariable* lhs);
     virtual llvm::Value* makeClassCheck(IREmitter& emitter, VAR* var, BoxedClass* c) {
         printf("makeClassCheck not defined for %s\n", debugName().c_str());
         abort();
@@ -277,6 +278,7 @@ public:
     virtual CompilerVariable* getPystonIter(IREmitter& emitter, const OpInfo& info) = 0;
     virtual CompilerVariable* binexp(IREmitter& emitter, const OpInfo& info, CompilerVariable* rhs,
                                      AST_TYPE::AST_TYPE op_type, BinExpType exp_type) = 0;
+    virtual CompilerVariable* contains(IREmitter& emitter, const OpInfo& info, CompilerVariable* lhs) = 0;
 
     virtual void serializeToFrame(std::vector<llvm::Value*>& stackmap_args) = 0;
 
@@ -369,6 +371,9 @@ public:
                              BinExpType exp_type) override {
         return type->binexp(emitter, info, this, rhs, op_type, exp_type);
     }
+    CompilerVariable* contains(IREmitter& emitter, const OpInfo& info, CompilerVariable* lhs) override {
+        return type->contains(emitter, info, this, lhs);
+    }
 
     llvm::Value* makeClassCheck(IREmitter& emitter, BoxedClass* cls) override {
         return type->makeClassCheck(emitter, this, cls);
@@ -391,6 +396,9 @@ public:
 // assert(value->getType() == type->llvmType());
 //}
 
+// Emit the test for whether one variable 'is' another one.
+ConcreteCompilerVariable* doIs(IREmitter& emitter, CompilerVariable* lhs, CompilerVariable* rhs, bool negate);
+
 ConcreteCompilerVariable* makeBool(bool);
 ConcreteCompilerVariable* makeInt(int64_t);
 ConcreteCompilerVariable* makeFloat(double);
@@ -421,6 +429,15 @@ CompilerVariable* _ValuedCompilerType<V>::getPystonIter(IREmitter& emitter, cons
     return r;
 }
 
+template <typename V>
+CompilerVariable* _ValuedCompilerType<V>::contains(IREmitter& emitter, const OpInfo& info, VAR* var,
+                                                   CompilerVariable* rhs) {
+    ConcreteCompilerVariable* converted = makeConverted(emitter, var, getBoxType());
+    auto r = UNKNOWN->contains(emitter, info, converted, rhs);
+    converted->decvref(emitter);
+    return r;
+}
+
 template <typename V>
 std::vector<CompilerVariable*> _ValuedCompilerType<V>::unpack(IREmitter& emitter, const OpInfo& info, VAR* var,
                                                               int num_into) {

src/codegen/irgen/irgenerator.cpp

@@ -766,6 +766,19 @@ private:
         assert(left);
         assert(right);
 
+        if (type == AST_TYPE::In || type == AST_TYPE::NotIn) {
+            CompilerVariable* r = right->contains(emitter, getOpInfoForNode(node, unw_info), left);
+            assert(r->getType() == BOOL);
+            if (type == AST_TYPE::NotIn) {
+                ConcreteCompilerVariable* converted = r->makeConverted(emitter, BOOL);
+                // TODO: would be faster to just do unboxBoolNegated
+                llvm::Value* raw = i1FromBool(emitter, converted);
+                raw = emitter.getBuilder()->CreateXor(raw, getConstantInt(1, g.i1));
+                r = boolFromI1(emitter, raw);
+            }
+            return r;
+        }
+
         return left->binexp(emitter, getOpInfoForNode(node, unw_info), right, type, exp_type);
     }
 
@@ -803,17 +816,7 @@ private:
         assert(right);
 
         if (node->ops[0] == AST_TYPE::Is || node->ops[0] == AST_TYPE::IsNot) {
-            // TODO: I think we can do better here and not force the types to box themselves
-
-            ConcreteCompilerVariable* converted_left = left->makeConverted(emitter, UNKNOWN);
-            ConcreteCompilerVariable* converted_right = right->makeConverted(emitter, UNKNOWN);
-            llvm::Value* cmp;
-            if (node->ops[0] == AST_TYPE::Is)
-                cmp = emitter.getBuilder()->CreateICmpEQ(converted_left->getValue(), converted_right->getValue());
-            else
-                cmp = emitter.getBuilder()->CreateICmpNE(converted_left->getValue(), converted_right->getValue());
-
-            return boolFromI1(emitter, cmp);
+            return doIs(emitter, left, right, node->ops[0] == AST_TYPE::IsNot);
         }
 
         CompilerVariable* rtn = _evalBinExp(node, left, right, node->ops[0], Compare, unw_info);

src/codegen/irgen/util.cpp

@@ -219,11 +219,11 @@ public:
                         if (!lookup_success) {
                             llvm::Constant* int_val
                                 = llvm::ConstantInt::get(g.i64, reinterpret_cast<uintptr_t>(addr), false);
-                            llvm::Constant* ptr_val = llvm::ConstantExpr::getIntToPtr(int_val, g.i8_ptr);
+                            llvm::Constant* ptr_val = llvm::ConstantExpr::getIntToPtr(int_val, g.i8);
                             ii->setArgOperand(i, ptr_val);
                             continue;
                         } else {
-                            ii->setArgOperand(i, module->getOrInsertGlobal(name, g.i8_ptr));
+                            ii->setArgOperand(i, module->getOrInsertGlobal(name, g.i8));
                         }
                     }
                 }

src/runtime/tuple.cpp

@@ -223,10 +223,8 @@ Box* tupleNonzero(BoxedTuple* self) {
 
 Box* tupleContains(BoxedTuple* self, Box* elt) {
     int size = self->size();
-    for (int i = 0; i < size; i++) {
-        Box* e = self->elts[i];
-
-        int r = PyObject_RichCompareBool(e, elt, Py_EQ);
+    for (Box* e : *self) {
+        int r = PyObject_RichCompareBool(elt, e, Py_EQ);
         if (r == -1)
             throwCAPIException();
 

test/tests/tuples.py

@@ -221,3 +221,6 @@ try:
     print (1, 3, 5, 3).index(2)
 except ValueError as e:
     print e
+
+n = float('nan')
+print n in (n, n)