Add compvars and rewriting support for GET_ITER nodes

If the type analysis successfully determines that the __iter__() and __hasnext__()
methods we can now replace them with direct calls to the destination.
Else we create a patchpoint and do a rewrite.

src/analysis/type_analysis.cpp

@@ -365,8 +365,9 @@ private:
                 return BOOL;
             case AST_LangPrimitive::LOCALS:
                 return DICT;
-            case AST_LangPrimitive::LANDINGPAD:
             case AST_LangPrimitive::GET_ITER:
+                return getType(node->args[0])->getPystonIterType();
+            case AST_LangPrimitive::LANDINGPAD:
             case AST_LangPrimitive::IMPORT_FROM:
             case AST_LangPrimitive::IMPORT_STAR:
             case AST_LangPrimitive::IMPORT_NAME:

src/codegen/compvars.cpp

@@ -34,6 +34,29 @@
 
 namespace pyston {
 
+static const std::string& iter_str = "__iter__";
+static const std::string& hasnext_str = "__hasnext__";
+
+CompilerType* CompilerType::getPystonIterType() {
+    if (hasattr(&iter_str) == Yes) {
+        CompilerType* iter_type = getattrType(&iter_str, true)->callType(ArgPassSpec(0), {}, NULL);
+        if (iter_type->hasattr(&hasnext_str) == Yes)
+            return iter_type;
+        // if iter_type->hasattr(&hasnext_str) == No we know this is going to be a BoxedIterWrapper
+        // we could optimize this case but it looks like this is very uncommon
+    }
+    return UNKNOWN;
+}
+
+CompilerType::Result CompilerType::hasattr(const std::string* attr) {
+    CompilerType* type = getattrType(attr, true);
+    if (type == UNKNOWN)
+        return Result::Maybe;
+    else if (type == UNDEF)
+        return Result::No;
+    return Result::Yes;
+}
+
 void ConcreteCompilerType::serializeToFrame(VAR* var, std::vector<llvm::Value*>& stackmap_args) {
 #ifndef NDEBUG
     if (llvmType() == g.i1) {
@@ -330,6 +353,57 @@ public:
         return new ConcreteCompilerVariable(UNKNOWN, rtn, true);
     }
 
+    CompilerVariable* getPystonIter(IREmitter& emitter, const OpInfo& info, ConcreteCompilerVariable* var) override {
+        CallattrFlags flags = {.cls_only = true, .null_on_nonexistent = true };
+        CompilerVariable* iter_call = var->callattr(emitter, info, &iter_str, flags, ArgPassSpec(0), {}, 0);
+        ConcreteCompilerVariable* converted_iter_call = iter_call->makeConverted(emitter, iter_call->getBoxType());
+
+        // If the type analysis could determine the iter type is a valid pyston iter (has 'hasnext') we are finished.
+        CompilerType* iter_type = var->getType()->getPystonIterType();
+        if (iter_type != UNKNOWN) {
+            iter_call->decvref(emitter);
+            return converted_iter_call;
+        }
+
+        // We don't know the type so we have to check at runtime if __iter__ is implemented
+        llvm::Value* cmp = emitter.getBuilder()->CreateICmpNE(converted_iter_call->getValue(),
+                                                              embedConstantPtr(0, g.llvm_value_type_ptr));
+
+        llvm::BasicBlock* bb_has_iter = emitter.createBasicBlock("has_iter");
+        bb_has_iter->moveAfter(emitter.currentBasicBlock());
+        llvm::BasicBlock* bb_no_iter = emitter.createBasicBlock("no_iter");
+        bb_no_iter->moveAfter(bb_has_iter);
+        llvm::BasicBlock* bb_join = emitter.createBasicBlock("join_after_getiter");
+        emitter.getBuilder()->CreateCondBr(cmp, bb_has_iter, bb_no_iter);
+
+        // var has __iter__()
+        emitter.setCurrentBasicBlock(bb_has_iter);
+        ICSetupInfo* pp = createGenericIC(info.getTypeRecorder(), true, 128);
+        llvm::Value* uncasted = emitter.createIC(pp, (void*)pyston::createBoxedIterWrapperIfNeeded,
+                                                 { converted_iter_call->getValue() }, info.unw_info);
+        llvm::Value* value_has_iter = emitter.getBuilder()->CreateIntToPtr(uncasted, g.llvm_value_type_ptr);
+        llvm::BasicBlock* value_has_iter_bb = emitter.currentBasicBlock();
+        emitter.getBuilder()->CreateBr(bb_join);
+
+        // var has no __iter__()
+        // TODO: we could create a patchpoint if this turns out to be hot
+        emitter.setCurrentBasicBlock(bb_no_iter);
+        llvm::Value* value_no_iter = emitter.createCall(info.unw_info, g.funcs.getiterHelper, var->getValue());
+        llvm::BasicBlock* value_no_iter_bb = emitter.currentBasicBlock();
+        emitter.getBuilder()->CreateBr(bb_join);
+
+        // join
+        emitter.setCurrentBasicBlock(bb_join);
+        auto phi = emitter.getBuilder()->CreatePHI(g.llvm_value_type_ptr, 2, "iter");
+        phi->addIncoming(value_has_iter, value_has_iter_bb);
+        phi->addIncoming(value_no_iter, value_no_iter_bb);
+
+        converted_iter_call->decvref(emitter);
+        iter_call->decvref(emitter);
+
+        return new ConcreteCompilerVariable(UNKNOWN, phi, true);
+    }
+
     CompilerVariable* binexp(IREmitter& emitter, const OpInfo& info, VAR* var, CompilerVariable* rhs,
                              AST_TYPE::AST_TYPE op_type, BinExpType exp_type) override {
         ConcreteCompilerVariable* converted_rhs = rhs->makeConverted(emitter, rhs->getBoxType());
@@ -1607,6 +1681,10 @@ public:
         return UNKNOWN->getitem(emitter, info, var, slice);
     }
 
+    CompilerVariable* getPystonIter(IREmitter& emitter, const OpInfo& info, VAR* var) override {
+        return UNKNOWN->getPystonIter(emitter, info, var);
+    }
+
     ConcreteCompilerVariable* len(IREmitter& emitter, const OpInfo& info, VAR* var) override {
         static const std::string attr("__len__");
         ConcreteCompilerVariable* called_constant

src/codegen/compvars.h

@@ -36,14 +36,19 @@ typedef llvm::SmallVector<uint64_t, 1> FrameVals;
 
 class CompilerType {
 public:
+    enum Result { Yes, No, Maybe };
+
     virtual ~CompilerType() {}
     virtual std::string debugName() = 0;
     virtual ConcreteCompilerType* getConcreteType() = 0;
     virtual ConcreteCompilerType* getBoxType() = 0;
     virtual bool canConvertTo(ConcreteCompilerType* other_type) = 0;
     virtual CompilerType* getattrType(const std::string* attr, bool cls_only) = 0;
+    virtual CompilerType* getPystonIterType();
+    virtual Result hasattr(const std::string* attr);
     virtual CompilerType* callType(ArgPassSpec argspec, const std::vector<CompilerType*>& arg_types,
                                    const std::vector<const std::string*>* keyword_names) = 0;
+
     virtual BoxedClass* guaranteedClass() = 0;
     virtual Box* deserializeFromFrame(const FrameVals& vals) = 0;
     virtual int numFrameArgs() = 0;
@@ -135,6 +140,7 @@ public:
         printf("getitem not defined for %s\n", debugName().c_str());
         abort();
     }
+    virtual CompilerVariable* getPystonIter(IREmitter& emitter, const OpInfo& info, VAR* var);
     virtual CompilerVariable* binexp(IREmitter& emitter, const OpInfo& info, VAR* var, CompilerVariable* rhs,
                                      AST_TYPE::AST_TYPE op_type, BinExpType exp_type) {
         printf("binexp not defined for %s\n", debugName().c_str());
@@ -263,6 +269,7 @@ public:
                                    const std::vector<const std::string*>* keyword_names) = 0;
     virtual ConcreteCompilerVariable* len(IREmitter& emitter, const OpInfo& info) = 0;
     virtual CompilerVariable* getitem(IREmitter& emitter, const OpInfo& info, CompilerVariable*) = 0;
+    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;
 
@@ -347,7 +354,9 @@ public:
     CompilerVariable* getitem(IREmitter& emitter, const OpInfo& info, CompilerVariable* slice) override {
         return type->getitem(emitter, info, this, slice);
     }
-
+    CompilerVariable* getPystonIter(IREmitter& emitter, const OpInfo& info) override {
+        return type->getPystonIter(emitter, info, this);
+    }
     CompilerVariable* binexp(IREmitter& emitter, const OpInfo& info, CompilerVariable* rhs, AST_TYPE::AST_TYPE op_type,
                              BinExpType exp_type) override {
         return type->binexp(emitter, info, this, rhs, op_type, exp_type);
@@ -392,6 +401,14 @@ CompilerType* makeFuncType(ConcreteCompilerType* rtn_type, const std::vector<Con
 ConcreteCompilerVariable* boolFromI1(IREmitter&, llvm::Value*);
 llvm::Value* i1FromBool(IREmitter&, ConcreteCompilerVariable*);
 
+template <typename V>
+CompilerVariable* _ValuedCompilerType<V>::getPystonIter(IREmitter& emitter, const OpInfo& info, VAR* var) {
+    ConcreteCompilerVariable* converted = makeConverted(emitter, var, getBoxType());
+    auto r = UNKNOWN->getPystonIter(emitter, info, converted);
+    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.h

@@ -68,8 +68,11 @@ public:
     virtual IRBuilder* getBuilder() = 0;
     virtual GCBuilder* getGC() = 0;
     virtual CompiledFunction* currentFunction() = 0;
+    virtual llvm::BasicBlock* currentBasicBlock() = 0;
     virtual llvm::BasicBlock* createBasicBlock(const char* name = "") = 0;
 
+    virtual void setCurrentBasicBlock(llvm::BasicBlock*) = 0;
+
     virtual llvm::Value* getScratch(int num_bytes) = 0;
     virtual void releaseScratch(llvm::Value*) = 0;
 

src/codegen/irgen/irgenerator.cpp

@@ -127,9 +127,6 @@ private:
         }
     }
 
-    llvm::BasicBlock* createBasicBlock(const char* name) override {
-        return llvm::BasicBlock::Create(g.context, name, irstate->getLLVMFunction());
-    }
 
     llvm::CallSite emitPatchpoint(llvm::Type* return_type, const ICSetupInfo* pp, llvm::Value* func,
                                   const std::vector<llvm::Value*>& args,
@@ -194,6 +191,16 @@ public:
     void releaseScratch(llvm::Value* scratch) override { assert(0); }
 
     CompiledFunction* currentFunction() override { return irstate->getCurFunction(); }
+    llvm::BasicBlock* currentBasicBlock() override { return curblock; }
+
+    void setCurrentBasicBlock(llvm::BasicBlock* bb) override {
+        curblock = bb;
+        getBuilder()->SetInsertPoint(curblock);
+    }
+
+    llvm::BasicBlock* createBasicBlock(const char* name) override {
+        return llvm::BasicBlock::Create(g.context, name, irstate->getLLVMFunction());
+    }
 
     llvm::Value* createCall(UnwindInfo unw_info, llvm::Value* callee, const std::vector<llvm::Value*>& args) override {
         if (ENABLE_FRAME_INTROSPECTION) {
@@ -489,24 +496,11 @@ private:
                 return rtn;
             }
             case AST_LangPrimitive::GET_ITER: {
-                // TODO if this is a type that has an __iter__, we could do way better than this, both in terms of
-                // function call overhead and resulting type information, if we went with that instead of the generic
-                // version.
-                // (ie we can inline getPystonIter here, whether mechanically with LLVM [would require adding more
-                // optimization passes to make it fast] or by-hand)
-                //
-                // TODO Move this behavior into to the type-specific section (compvars.cpp)?
-                emitter.getBuilder();
                 assert(node->args.size() == 1);
                 CompilerVariable* obj = evalExpr(node->args[0], unw_info);
-
-                ConcreteCompilerVariable* converted_obj = obj->makeConverted(emitter, obj->getBoxType());
+                auto rtn = obj->getPystonIter(emitter, getOpInfoForNode(node, unw_info));
                 obj->decvref(emitter);
-
-                llvm::Value* v = emitter.createCall(unw_info, g.funcs.getPystonIter, converted_obj->getValue());
-                assert(v->getType() == g.llvm_value_type_ptr);
-
-                return new ConcreteCompilerVariable(UNKNOWN, v, true);
+                return rtn;
             }
             case AST_LangPrimitive::IMPORT_FROM: {
                 assert(node->args.size() == 2);

src/codegen/runtime_hooks.cpp

@@ -207,7 +207,7 @@ void initGlobalFuncs(GlobalState& g) {
     GET(str);
     GET(isinstance);
     GET(yield);
-    GET(getPystonIter);
+    GET(getiterHelper);
 
     GET(unpackIntoArray);
     GET(raiseAttributeError);

src/codegen/runtime_hooks.h

@@ -37,7 +37,7 @@ struct GlobalFuncs {
         *createUserClass, *createClosure, *createGenerator, *createLong, *createSet, *createPureImaginary;
     llvm::Value* getattr, *setattr, *delattr, *delitem, *delGlobal, *nonzero, *binop, *compare, *augbinop, *unboxedLen,
         *getitem, *getclsattr, *getGlobal, *setitem, *unaryop, *import, *importFrom, *importStar, *repr, *str,
-        *isinstance, *yield, *getPystonIter;
+        *isinstance, *yield, *getiterHelper;
 
     llvm::Value* unpackIntoArray, *raiseAttributeError, *raiseAttributeErrorStr, *raiseNotIterableError,
         *assertNameDefined, *assertFail;

src/runtime/inline/link_forcer.cpp

@@ -88,7 +88,7 @@ void force() {
     FORCE(str);
     FORCE(isinstance);
     FORCE(yield);
-    FORCE(getPystonIter);
+    FORCE(getiterHelper);
 
     FORCE(unpackIntoArray);
     FORCE(raiseAttributeError);

src/runtime/objmodel.cpp

@@ -3461,6 +3461,41 @@ extern "C" void delattr(Box* obj, const char* attr) {
     delattr_internal(obj, attr, true, NULL);
 }
 
+extern "C" Box* createBoxedIterWrapper(Box* o) {
+    return new BoxedIterWrapper(o);
+}
+
+extern "C" Box* createBoxedIterWrapperIfNeeded(Box* o) {
+    std::unique_ptr<Rewriter> rewriter(Rewriter::createRewriter(
+        __builtin_extract_return_addr(__builtin_return_address(0)), 1, "createBoxedIterWrapperIfNeeded"));
+
+    if (rewriter.get()) {
+        RewriterVar* r_o = rewriter->getArg(0);
+        RewriterVar* r_cls = r_o->getAttr(BOX_CLS_OFFSET);
+        GetattrRewriteArgs rewrite_args(rewriter.get(), r_cls, rewriter->getReturnDestination());
+        Box* r = typeLookup(o->cls, hasnext_str, &rewrite_args);
+        if (!rewrite_args.out_success) {
+            rewriter.reset(NULL);
+        } else if (r) {
+            rewrite_args.out_rtn->addGuard((uint64_t)r);
+            if (rewrite_args.out_success) {
+                rewriter->commitReturning(r_o);
+                return o;
+            }
+        } else if (!r) {
+            RewriterVar* var = rewriter.get()->call(true, (void*)createBoxedIterWrapper, rewriter->getArg(0));
+            if (rewrite_args.out_success) {
+                rewriter->commitReturning(var);
+                return createBoxedIterWrapper(o);
+            }
+        }
+    }
+
+    if (typeLookup(o->cls, hasnext_str, NULL) == NULL)
+        return new BoxedIterWrapper(o);
+    return o;
+}
+
 extern "C" Box* getPystonIter(Box* o) {
     Box* r = getiter(o);
     if (typeLookup(r->cls, hasnext_str, NULL) == NULL)
@@ -3468,17 +3503,18 @@ extern "C" Box* getPystonIter(Box* o) {
     return r;
 }
 
+extern "C" Box* getiterHelper(Box* o) {
+    if (typeLookup(o->cls, getitem_str, NULL))
+        return new BoxedSeqIter(o, 0);
+    raiseExcHelper(TypeError, "'%s' object is not iterable", getTypeName(o));
+}
+
 Box* getiter(Box* o) {
     // TODO add rewriting to this?  probably want to try to avoid this path though
     Box* r = callattrInternal0(o, &iter_str, LookupScope::CLASS_ONLY, NULL, ArgPassSpec(0));
     if (r)
         return r;
-
-    if (typeLookup(o->cls, getitem_str, NULL)) {
-        return new BoxedSeqIter(o, 0);
-    }
-
-    raiseExcHelper(TypeError, "'%s' object is not iterable", getTypeName(o));
+    return getiterHelper(o);
 }
 
 llvm::iterator_range<BoxIterator> Box::pyElements() {

src/runtime/objmodel.h

@@ -87,6 +87,8 @@ extern "C" BoxedClosure* createClosure(BoxedClosure* parent_closure);
 
 Box* getiter(Box* o);
 extern "C" Box* getPystonIter(Box* o);
+extern "C" Box* getiterHelper(Box* o);
+extern "C" Box* createBoxedIterWrapperIfNeeded(Box* o);
 
 extern "C" void dump(void* p);
 

src/runtime/types.h

@@ -226,7 +226,7 @@ public:
     // that we can't rely on for extension classes.
     bool is_pyston_class;
 
-    bool hasGenericGetattr() { return tp_getattr != NULL; }
+    bool hasGenericGetattr() { return tp_getattr == NULL; }
 
     void freeze();
 

test/tests/xrange.py

@@ -2,7 +2,7 @@
 #
 # statcheck: "-O" in EXTRA_JIT_ARGS or 'slowpath_getclsattr' in stats or 'slowpath_callattr' in stats
 # statcheck: stats.get('slowpath_getclsattr', 0) <= 20
-# statcheck: stats.get('slowpath_callattr', 0) <= 20
+# statcheck: stats.get('slowpath_callattr', 0) <= 22
 
 for i in xrange(1000):
     print i