Optimize seq_iter protocol

- special-case classes that we know have a fixed length (ex str)
- do patchpoints for getIterHelper

src/codegen/compvars.cpp

@@ -438,8 +438,13 @@ public:
         // 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());
+
+        ICSetupInfo* pp2 = createGenericIC(info.getTypeRecorder(), true, 128);
+        llvm::Instruction* value_no_iter
+            = emitter.createIC(pp2, (void*)getiterHelper, { var->getValue() }, info.unw_info);
+        value_no_iter = createAfter<llvm::IntToPtrInst>(value_no_iter, value_no_iter, g.llvm_value_type_ptr, "");
         emitter.setType(value_no_iter, RefType::OWNED);
+
         llvm::BasicBlock* value_no_iter_bb = emitter.currentBasicBlock();
         auto no_iter_terminator = emitter.getBuilder()->CreateBr(bb_join);
 

src/core/types.h

@@ -682,6 +682,7 @@ extern HiddenClass* root_hcls;
 struct SetattrRewriteArgs;
 struct GetattrRewriteArgs;
 struct DelattrRewriteArgs;
+struct UnaryopRewriteArgs;
 
 // Helper function around PyString_InternFromString:
 BoxedString* internStringImmortal(llvm::StringRef s) noexcept;

src/runtime/iterobject.cpp

@@ -48,6 +48,19 @@ static Box* seqiterHasnext_capi(Box* s) noexcept {
         Py_RETURN_FALSE;
     }
 
+    if (self->len != -1) {
+        if (self->idx >= self->len)
+            Py_RETURN_FALSE;
+        assert(!self->next);
+        if (self->b->cls == str_cls)
+            self->next = incref(characters[static_cast<BoxedString*>(self->b)->s()[self->idx] & UCHAR_MAX]);
+        else
+            self->next = PySequence_GetItem(self->b, self->idx);
+        assert(self->next);
+        self->idx++;
+        Py_RETURN_TRUE;
+    }
+
     Box* next = PySequence_GetItem(self->b, self->idx);
     if (!next) {
         if (PyErr_ExceptionMatches(IndexError) || PyErr_ExceptionMatches(StopIteration)) {
@@ -199,7 +212,7 @@ void setupIter() {
 
     seqiter_cls->giveAttr("next", new BoxedFunction(FunctionMetadata::create((void*)seqiterNext, UNKNOWN, 1)));
     seqiter_cls->giveAttr("__hasnext__",
-                          new BoxedFunction(FunctionMetadata::create((void*)seqiterHasnext, BOXED_BOOL, 1)));
+                          new BoxedFunction(FunctionMetadata::create((void*)seqiterHasnext_capi, BOXED_BOOL, 1, ParamNames::empty(), CAPI)));
     seqiter_cls->giveAttr("__iter__", new BoxedFunction(FunctionMetadata::create((void*)seqiterIter, UNKNOWN, 1)));
 
     seqiter_cls->freeze();

src/runtime/iterobject.h

@@ -33,9 +33,22 @@ public:
     int64_t idx;
     Box* next;
 
-    BoxedSeqIter(Box* b, int64_t start) : b(b), idx(start), next(NULL) { Py_INCREF(b); }
+    // Pyston change:
+    // For types that allow it, this class will do the more efficient length-based
+    // iteration, storing the length here.  Otherwise len is -1.
+    int64_t len;
+
+    BoxedSeqIter(Box* b, int64_t start) : b(b), idx(start), next(NULL) {
+        Py_INCREF(b);
+
+        if (b->cls == str_cls) {
+            len = static_cast<BoxedString*>(b)->size();
+        } else {
+            len = -1;
+        }
+    }
 
-    DEFAULT_CLASS(seqiter_cls);
+    DEFAULT_CLASS_SIMPLE(seqiter_cls, true);
 
     static void dealloc(BoxedSeqIter* o) noexcept {
         PyObject_GC_UnTrack(o);

src/runtime/objmodel.cpp

@@ -3466,7 +3466,7 @@ extern "C" BoxedInt* hash(Box* obj) {
 }
 
 template <ExceptionStyle S, Rewritable rewritable>
-BoxedInt* lenInternal(Box* obj, LenRewriteArgs* rewrite_args) noexcept(S == CAPI) {
+BoxedInt* lenInternal(Box* obj, UnaryopRewriteArgs* rewrite_args) noexcept(S == CAPI) {
     if (rewritable == NOT_REWRITABLE) {
         assert(!rewrite_args);
         rewrite_args = NULL;
@@ -3594,10 +3594,10 @@ BoxedInt* lenInternal(Box* obj, LenRewriteArgs* rewrite_args) noexcept(S == CAPI
 }
 
 // force template instantiation:
-template BoxedInt* lenInternal<CAPI, REWRITABLE>(Box*, LenRewriteArgs*);
-template BoxedInt* lenInternal<CXX, REWRITABLE>(Box*, LenRewriteArgs*);
-template BoxedInt* lenInternal<CAPI, NOT_REWRITABLE>(Box*, LenRewriteArgs*);
-template BoxedInt* lenInternal<CXX, NOT_REWRITABLE>(Box*, LenRewriteArgs*);
+template BoxedInt* lenInternal<CAPI, REWRITABLE>(Box*, UnaryopRewriteArgs*);
+template BoxedInt* lenInternal<CXX, REWRITABLE>(Box*, UnaryopRewriteArgs*);
+template BoxedInt* lenInternal<CAPI, NOT_REWRITABLE>(Box*, UnaryopRewriteArgs*);
+template BoxedInt* lenInternal<CXX, NOT_REWRITABLE>(Box*, UnaryopRewriteArgs*);
 
 Box* lenCallInternal(BoxedFunctionBase* func, CallRewriteArgs* rewrite_args, ArgPassSpec argspec, Box* arg1, Box* arg2,
                      Box* arg3, Box** args, const std::vector<BoxedString*>* keyword_names) {
@@ -3605,7 +3605,7 @@ Box* lenCallInternal(BoxedFunctionBase* func, CallRewriteArgs* rewrite_args, Arg
         return callFunc<CXX>(func, rewrite_args, argspec, arg1, arg2, arg3, args, keyword_names);
 
     if (rewrite_args) {
-        LenRewriteArgs lrewrite_args(rewrite_args->rewriter, rewrite_args->arg1, rewrite_args->destination);
+        UnaryopRewriteArgs lrewrite_args(rewrite_args->rewriter, rewrite_args->arg1, rewrite_args->destination);
         Box* rtn = lenInternal<CXX, REWRITABLE>(arg1, &lrewrite_args);
         if (!lrewrite_args.out_success) {
             rewrite_args = 0;
@@ -3640,7 +3640,7 @@ extern "C" i64 unboxedLen(Box* obj) {
     RewriterVar* r_boxed = NULL;
     if (rewriter.get()) {
         // rewriter->trap();
-        LenRewriteArgs rewrite_args(rewriter.get(), rewriter->getArg(0), rewriter->getReturnDestination());
+        UnaryopRewriteArgs rewrite_args(rewriter.get(), rewriter->getArg(0), rewriter->getReturnDestination());
         lobj = lenInternal<CXX, REWRITABLE>(obj, &rewrite_args);
 
         if (!rewrite_args.out_success) {
@@ -6804,14 +6804,50 @@ extern "C" Box* getPystonIter(Box* o) {
     return r;
 }
 
-extern "C" Box* getiterHelper(Box* o) {
-    if (PySequence_Check(o))
+extern "C" Box* getiterHelperInternal(Box* o, UnaryopRewriteArgs* rewrite_args) {
+    if (rewrite_args)
+        assert(o->cls->is_constant);
+
+    if (PySequence_Check(o)) {
+        class Helper {
+        public:
+            static Box* call(Box* o) { return new BoxedSeqIter(o, 0); }
+        };
+
+        if (rewrite_args) {
+            rewrite_args->out_rtn
+                = rewrite_args->rewriter->call(false, (void*)&Helper::call, rewrite_args->obj)->setType(RefType::OWNED);
+            rewrite_args->out_success = true;
+        }
         return new BoxedSeqIter(o, 0);
+    }
+
     raiseExcHelper(TypeError, "'%s' object is not iterable", getTypeName(o));
 }
 
+extern "C" Box* getiterHelper(Box* o) {
+    std::unique_ptr<Rewriter> rewriter(nullptr);
+
+    if (o->cls->is_constant)
+        rewriter.reset(
+            Rewriter::createRewriter(__builtin_extract_return_addr(__builtin_return_address(0)), 1, "getPystonIter"));
+
+    if (rewriter.get()) {
+        UnaryopRewriteArgs rewrite_args(rewriter.get(), rewriter->getArg(0)->setType(RefType::BORROWED),
+                                        rewriter->getReturnDestination());
+
+        Box* r = getiterHelperInternal(o, &rewrite_args);
+        if (rewrite_args.out_success) {
+            RewriterVar* r_rtn = rewrite_args.out_rtn;
+            rewriter->commitReturning(r_rtn);
+        }
+        return r;
+    } else {
+        return getiterHelperInternal(o, NULL);
+    }
+}
+
 Box* getiter(Box* o) {
-    // TODO add rewriting to this?  probably want to try to avoid this path though
     BoxedClass* type = o->cls;
     Box* r = NULL;
     if (PyType_HasFeature(type, Py_TPFLAGS_HAVE_ITER) && type->tp_iter != slot_tp_iter && type->tp_iter) {
@@ -6828,7 +6864,7 @@ Box* getiter(Box* o) {
         }
         return r;
     }
-    return getiterHelper(o);
+    return getiterHelperInternal(o, NULL);
 }
 
 void assertValidSlotIdentifier(Box* s) {

src/runtime/objmodel.h

@@ -110,7 +110,7 @@ extern "C" BoxedClosure* createClosure(BoxedClosure* parent_closure, size_t size
 
 Box* getiter(Box* o);
 extern "C" Box* getPystonIter(Box* o);
-extern "C" Box* getiterHelper(Box* o);
+extern "C" Box* getiterHelper(Box* o) __attribute__((noinline));
 extern "C" Box* createBoxedIterWrapperIfNeeded(Box* o) __attribute__((noinline));
 
 struct SetattrRewriteArgs;
@@ -133,9 +133,8 @@ template <ExceptionStyle S> inline Box* getitemInternal(Box* target, Box* slice)
     return getitemInternal<S, NOT_REWRITABLE>(target, slice, NULL);
 }
 
-struct LenRewriteArgs;
 template <ExceptionStyle S, Rewritable rewritable>
-BoxedInt* lenInternal(Box* obj, LenRewriteArgs* rewrite_args) noexcept(S == CAPI);
+BoxedInt* lenInternal(Box* obj, UnaryopRewriteArgs* rewrite_args) noexcept(S == CAPI);
 Box* lenCallInternal(BoxedFunctionBase* f, CallRewriteArgs* rewrite_args, ArgPassSpec argspec, Box* arg1, Box* arg2,
                      Box* arg3, Box** args, const std::vector<BoxedString*>* keyword_names);
 

src/runtime/rewrite_args.h

@@ -194,7 +194,7 @@ struct DelattrRewriteArgs {
     DelattrRewriteArgs(Rewriter* rewriter, RewriterVar* obj) : rewriter(rewriter), obj(obj), out_success(false) {}
 };
 
-struct LenRewriteArgs {
+struct UnaryopRewriteArgs {
     Rewriter* rewriter;
     RewriterVar* obj;
     Location destination;
@@ -202,7 +202,7 @@ struct LenRewriteArgs {
     bool out_success;
     RewriterVar* out_rtn;
 
-    LenRewriteArgs(Rewriter* rewriter, RewriterVar* obj, Location destination)
+    UnaryopRewriteArgs(Rewriter* rewriter, RewriterVar* obj, Location destination)
         : rewriter(rewriter), obj(obj), destination(destination), out_success(false), out_rtn(NULL) {}
 };