Merge pull request #318 from undingen/getiter2

Add compvars and rewriting support for GET_ITER nodes 

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.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;
@@ -112,7 +117,7 @@ public:
     }
 
     virtual CompilerVariable* callattr(IREmitter& emitter, const OpInfo& info, VAR* var, const std::string* attr,
-                                       bool clsonly, struct ArgPassSpec argspec,
+                                       CallattrFlags flags, struct ArgPassSpec argspec,
                                        const std::vector<CompilerVariable*>& args,
                                        const std::vector<const std::string*>* keyword_names) {
         printf("callattr not defined for %s\n", debugName().c_str());
@@ -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());
@@ -254,14 +260,16 @@ public:
         = 0;
     virtual void setattr(IREmitter& emitter, const OpInfo& info, const std::string* attr, CompilerVariable* v) = 0;
     virtual void delattr(IREmitter& emitter, const OpInfo& info, const std::string* attr) = 0;
-    virtual CompilerVariable* callattr(IREmitter& emitter, const OpInfo& info, const std::string* attr, bool clsonly,
-                                       struct ArgPassSpec argspec, const std::vector<CompilerVariable*>& args,
+    virtual CompilerVariable* callattr(IREmitter& emitter, const OpInfo& info, const std::string* attr,
+                                       CallattrFlags flags, struct ArgPassSpec argspec,
+                                       const std::vector<CompilerVariable*>& args,
                                        const std::vector<const std::string*>* keyword_names) = 0;
     virtual CompilerVariable* call(IREmitter& emitter, const OpInfo& info, struct ArgPassSpec argspec,
                                    const std::vector<CompilerVariable*>& args,
                                    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;
 
@@ -330,10 +338,10 @@ public:
         type->delattr(emitter, info, this, attr);
     }
 
-    CompilerVariable* callattr(IREmitter& emitter, const OpInfo& info, const std::string* attr, bool clsonly,
+    CompilerVariable* callattr(IREmitter& emitter, const OpInfo& info, const std::string* attr, CallattrFlags flags,
                                struct ArgPassSpec argspec, const std::vector<CompilerVariable*>& args,
                                const std::vector<const std::string*>* keyword_names) override {
-        return type->callattr(emitter, info, this, attr, clsonly, argspec, args, keyword_names);
+        return type->callattr(emitter, info, this, attr, flags, argspec, args, keyword_names);
     }
     CompilerVariable* call(IREmitter& emitter, const OpInfo& info, struct ArgPassSpec argspec,
                            const std::vector<CompilerVariable*>& args,
@@ -346,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 +402,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) {
@@ -457,8 +464,9 @@ private:
                         ConcreteCompilerVariable* converted = p.second->makeConverted(emitter, p.second->getBoxType());
 
                         // TODO super dumb that it reallocates the name again
+                        CallattrFlags flags = {.cls_only = true, .null_on_nonexistent = false };
                         CompilerVariable* _r
-                            = rtn->callattr(emitter, getEmptyOpInfo(unw_info), &setitem_str, true, ArgPassSpec(2),
+                            = rtn->callattr(emitter, getEmptyOpInfo(unw_info), &setitem_str, flags, ArgPassSpec(2),
                                             { makeStr(new std::string(p.first.str())), converted }, NULL);
                         converted->decvref(emitter);
                         _r->decvref(emitter);
@@ -474,8 +482,9 @@ private:
                         emitter.getBuilder()->SetInsertPoint(was_defined);
                         ConcreteCompilerVariable* converted = p.second->makeConverted(emitter, p.second->getBoxType());
                         // TODO super dumb that it reallocates the name again
+                        CallattrFlags flags = {.cls_only = true, .null_on_nonexistent = false };
                         CompilerVariable* _r
-                            = rtn->callattr(emitter, getEmptyOpInfo(unw_info), &setitem_str, true, ArgPassSpec(2),
+                            = rtn->callattr(emitter, getEmptyOpInfo(unw_info), &setitem_str, flags, ArgPassSpec(2),
                                             { makeStr(new std::string(p.first.str())), converted }, NULL);
                         converted->decvref(emitter);
                         _r->decvref(emitter);
@@ -487,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);
@@ -748,8 +744,8 @@ private:
 
         CompilerVariable* rtn;
         if (is_callattr) {
-            rtn = func->callattr(emitter, getOpInfoForNode(node, unw_info), attr, callattr_clsonly, argspec, args,
-                                 keyword_names);
+            CallattrFlags flags = {.cls_only = callattr_clsonly, .null_on_nonexistent = false };
+            rtn = func->callattr(emitter, getOpInfoForNode(node, unw_info), attr, flags, argspec, args, keyword_names);
         } else {
             rtn = func->call(emitter, getOpInfoForNode(node, unw_info), argspec, args, keyword_names);
         }
@@ -976,8 +972,8 @@ private:
 
         for (int i = 0; i < node->elts.size(); i++) {
             CompilerVariable* elt = elts[i];
-
-            CompilerVariable* r = rtn->callattr(emitter, getOpInfoForNode(node, unw_info), &add_str, true,
+            CallattrFlags flags = {.cls_only = true, .null_on_nonexistent = false };
+            CompilerVariable* r = rtn->callattr(emitter, getOpInfoForNode(node, unw_info), &add_str, flags,
                                                 ArgPassSpec(1), { elt }, NULL);
             r->decvref(emitter);
             elt->decvref(emitter);
@@ -1648,8 +1644,8 @@ private:
 
             curblock = ss_block;
             emitter.getBuilder()->SetInsertPoint(ss_block);
-
-            auto r = dest->callattr(emitter, getOpInfoForNode(node, unw_info), &write_str, false, ArgPassSpec(1),
+            CallattrFlags flags = {.cls_only = false, .null_on_nonexistent = false };
+            auto r = dest->callattr(emitter, getOpInfoForNode(node, unw_info), &write_str, flags, ArgPassSpec(1),
                                     { makeStr(&space_str) }, NULL);
             r->decvref(emitter);
 
@@ -1662,7 +1658,7 @@ private:
             llvm::Value* v = emitter.createCall(unw_info, g.funcs.str, converted->getValue());
             v = emitter.getBuilder()->CreateBitCast(v, g.llvm_value_type_ptr);
             auto s = new ConcreteCompilerVariable(STR, v, true);
-            r = dest->callattr(emitter, getOpInfoForNode(node, unw_info), &write_str, false, ArgPassSpec(1), { s },
+            r = dest->callattr(emitter, getOpInfoForNode(node, unw_info), &write_str, flags, ArgPassSpec(1), { s },
                                NULL);
             s->decvref(emitter);
             r->decvref(emitter);
@@ -1670,7 +1666,8 @@ private:
         }
 
         if (node->nl) {
-            auto r = dest->callattr(emitter, getOpInfoForNode(node, unw_info), &write_str, false, ArgPassSpec(1),
+            CallattrFlags flags = {.cls_only = false, .null_on_nonexistent = false };
+            auto r = dest->callattr(emitter, getOpInfoForNode(node, unw_info), &write_str, flags, ArgPassSpec(1),
                                     { makeStr(&newline_str) }, NULL);
             r->decvref(emitter);
 

src/codegen/runtime_hooks.cpp

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

src/codegen/runtime_hooks.h

@@ -38,7 +38,7 @@ struct GlobalFuncs {
         *decodeUTF8StringPtr;
     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/core/types.h

@@ -557,6 +557,11 @@ struct FrameInfo {
 
     FrameInfo(ExcInfo exc) : exc(exc) {}
 };
+
+struct CallattrFlags {
+    bool cls_only : 1;
+    bool null_on_nonexistent : 1;
+};
 }
 
 #endif

src/runtime/inline/link_forcer.cpp

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

src/runtime/objmodel.cpp

@@ -3471,6 +3471,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)
@@ -3478,17 +3513,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

@@ -53,10 +53,6 @@ extern "C" void setattr(Box* obj, const char* attr, Box* attr_val);
 extern "C" void delattr(Box* obj, const char* attr);
 extern "C" bool nonzero(Box* obj);
 extern "C" Box* runtimeCall(Box*, ArgPassSpec, Box*, Box*, Box*, Box**, const std::vector<const std::string*>*);
-struct CallattrFlags {
-    bool cls_only : 1;
-    bool null_on_nonexistent : 1;
-};
 extern "C" Box* callattr(Box*, const std::string*, CallattrFlags, ArgPassSpec, Box*, Box*, Box*, Box**,
                          const std::vector<const std::string*>*);
 extern "C" BoxedString* str(Box* obj);
@@ -92,6 +88,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

@@ -228,7 +228,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