Generators: remove single instance restriction and handle generators inside closures

src/analysis/scoping_analysis.cpp

@@ -24,6 +24,8 @@ class YieldVisitor : public NoopASTVisitor {
 public:
     YieldVisitor() : containsYield(false) {}
 
+    virtual bool visit_functiondef(AST_FunctionDef*) { return true; }
+
     virtual bool visit_yield(AST_Yield*) {
         containsYield = true;
         return true;
@@ -34,7 +36,16 @@ public:
 
 bool containsYield(AST* ast) {
     YieldVisitor visitor;
-    ast->accept(&visitor);
+    if (ast->type == AST_TYPE::FunctionDef) {
+        AST_FunctionDef* funcDef = static_cast<AST_FunctionDef*>(ast);
+        for (auto& e : funcDef->body) {
+            e->accept(&visitor);
+            if (visitor.containsYield)
+                return true;
+        }
+    } else {
+        ast->accept(&visitor);
+    }
     return visitor.containsYield;
 }
 
@@ -387,9 +398,12 @@ void ScopingAnalysis::processNameUsages(ScopingAnalysis::NameUsageMap* usages) {
         switch (node->type) {
             case AST_TYPE::ClassDef:
             case AST_TYPE::FunctionDef:
-            case AST_TYPE::Lambda:
-                this->scopes[node] = new ScopeInfoBase(parent_info, usage);
+            case AST_TYPE::Lambda: {
+                ScopeInfoBase* scopInfo = new ScopeInfoBase(parent_info, usage);
+                scopInfo->setTakesGenerator(containsYield(node));
+                this->scopes[node] = scopInfo;
                 break;
+            }
             default:
                 RELEASE_ASSERT(0, "%d", usage->node->type);
                 break;

src/codegen/compvars.cpp

@@ -518,8 +518,8 @@ ConcreteCompilerVariable* UnknownType::nonzero(IREmitter& emitter, const OpInfo&
     return new ConcreteCompilerVariable(BOOL, rtn_val, true);
 }
 
-CompilerVariable* makeFunction(IREmitter& emitter, CLFunction* f, CompilerVariable* closure,
-                               CompilerVariable* generator, const std::vector<ConcreteCompilerVariable*>& defaults) {
+CompilerVariable* makeFunction(IREmitter& emitter, CLFunction* f, CompilerVariable* closure, bool isGenerator,
+                               const std::vector<ConcreteCompilerVariable*>& defaults) {
     // Unlike the CLFunction*, which can be shared between recompilations, the Box* around it
     // should be created anew every time the functiondef is encountered
 
@@ -532,18 +532,6 @@ CompilerVariable* makeFunction(IREmitter& emitter, CLFunction* f, CompilerVariab
         closure_v = embedConstantPtr(nullptr, g.llvm_closure_type_ptr);
     }
 
-    llvm::Value* generator_v;
-    ConcreteCompilerVariable* convertedGenerator = NULL;
-
-    if (generator && generator != (CompilerVariable*)1) {
-        convertedGenerator = generator->makeConverted(emitter, generator->getConcreteType());
-        generator_v = convertedGenerator->getValue();
-        // ugly hack to allow to pass a BoxedFunction* instead of a BoxedGenerator*
-        generator_v = emitter.getBuilder()->CreateBitCast(generator_v, g.llvm_generator_type_ptr);
-    } else {
-        generator_v = embedConstantPtr(nullptr, g.llvm_generator_type_ptr);
-    }
-
     llvm::Value* scratch;
     if (defaults.size()) {
         scratch = emitter.getScratch(defaults.size() * sizeof(Box*));
@@ -559,15 +547,15 @@ CompilerVariable* makeFunction(IREmitter& emitter, CLFunction* f, CompilerVariab
         scratch = embedConstantPtr(nullptr, g.llvm_value_type_ptr_ptr);
     }
 
+    llvm::Value* isGenerator_v = llvm::ConstantInt::get(g.i1, isGenerator, false);
+
     llvm::Value* boxed = emitter.getBuilder()->CreateCall(
         g.funcs.boxCLFunction,
-        std::vector<llvm::Value*>{ embedConstantPtr(f, g.llvm_clfunction_type_ptr), closure_v, generator_v, scratch,
+        std::vector<llvm::Value*>{ embedConstantPtr(f, g.llvm_clfunction_type_ptr), closure_v, isGenerator_v, scratch,
                                    getConstantInt(defaults.size(), g.i64) });
 
     if (convertedClosure)
         convertedClosure->decvref(emitter);
-    if (convertedGenerator)
-        convertedGenerator->decvref(emitter);
     return new ConcreteCompilerVariable(typeFromClass(function_cls), boxed, true);
 }
 

src/codegen/compvars.h

@@ -316,7 +316,7 @@ ConcreteCompilerVariable* makeInt(int64_t);
 ConcreteCompilerVariable* makeFloat(double);
 ConcreteCompilerVariable* makeBool(bool);
 CompilerVariable* makeStr(std::string*);
-CompilerVariable* makeFunction(IREmitter& emitter, CLFunction*, CompilerVariable* closure, CompilerVariable* generator,
+CompilerVariable* makeFunction(IREmitter& emitter, CLFunction*, CompilerVariable* closure, bool isGenerator,
                                const std::vector<ConcreteCompilerVariable*>& defaults);
 ConcreteCompilerVariable* undefVariable();
 CompilerVariable* makeTuple(const std::vector<CompilerVariable*>& elts);

src/codegen/irgen/irgenerator.cpp

@@ -1474,7 +1474,7 @@ private:
         // one reason to do this is to pass the closure through if necessary,
         // but since the classdef can't create its own closure, shouldn't need to explicitly
         // create that scope to pass the closure through.
-        CompilerVariable* func = makeFunction(emitter, cl, created_closure, 0, {});
+        CompilerVariable* func = makeFunction(emitter, cl, created_closure, false, {});
 
         CompilerVariable* attr_dict = func->call(emitter, getEmptyOpInfo(exc_info), ArgPassSpec(0), {}, NULL);
 
@@ -1575,16 +1575,7 @@ private:
             assert(created_closure);
         }
 
-        CompilerVariable* func = makeFunction(emitter, cl, created_closure,
-                                              (ConcreteCompilerVariable*)scope_info->takesGenerator(), defaults);
-
-        if (scope_info->takesGenerator()) {
-            ConcreteCompilerVariable* converted = func->makeConverted(emitter, func->getBoxType());
-            CLFunction* clFunc = boxRTFunction((void*)createGenerator, UNKNOWN, args->args.size(),
-                                               args->defaults.size(), args->vararg.size(), args->kwarg.size());
-            func = makeFunction(emitter, clFunc, NULL, (CompilerVariable*)converted, defaults);
-            converted->decvref(emitter);
-        }
+        CompilerVariable* func = makeFunction(emitter, cl, created_closure, scope_info->takesGenerator(), defaults);
 
         for (auto d : defaults) {
             d->decvref(emitter);
@@ -1915,7 +1906,7 @@ private:
         llvm::BasicBlock* target = entry_blocks[node->target];
 
         if (ENABLE_OSR && node->target->idx < myblock->idx && irstate->getEffortLevel() < EffortLevel::MAXIMAL
-            && !containsYield(irstate->getSourceInfo()->ast)) {
+            && !irstate->getScopeInfo()->takesGenerator()) {
             assert(node->target->predecessors.size() > 1);
             doOSRExit(target, node);
         } else {

src/runtime/generator.cpp

@@ -20,6 +20,7 @@
 #include <ucontext.h>
 
 #include "codegen/compvars.h"
+#include "codegen/llvm_interpreter.h"
 #include "core/ast.h"
 #include "core/common.h"
 #include "core/stats.h"
@@ -33,23 +34,24 @@
 namespace pyston {
 
 
-static void generatorEntry(BoxedGenerator* self) {
-    assert(self->cls == generator_cls);
-    assert(self->function->cls == function_cls);
+static void generatorEntry(BoxedGenerator* g) {
+    assert(g->cls == generator_cls);
+    assert(g->function->cls == function_cls);
 
     try {
         // call body of the generator
-        ArgPassSpec argPassSpec(self->function->f->num_args, 0, self->function->f->takes_varargs,
-                                self->function->f->takes_kwargs);
-        runtimeCall(self->function, argPassSpec, self->arg1, self->arg2, self->arg3, self->args, 0);
+        BoxedFunction* func = g->function;
+
+        Box** args = g->args ? &g->args->elts[0] : nullptr;
+        callCLFunc(func->f, nullptr, func->f->num_args, func->closure, g, g->arg1, g->arg2, g->arg3, args);
     } catch (Box* e) {
         // unhandled exception: propagate the exception to the caller
-        self->exception = e;
+        g->exception = e;
     }
 
     // we returned from the body of the generator. next/send/throw will notify the caller
-    self->entryExited = true;
-    swapcontext(&self->context, &self->returnContext);
+    g->entryExited = true;
+    swapcontext(&g->context, &g->returnContext);
 }
 
 Box* generatorIter(Box* s) {
@@ -126,13 +128,18 @@ extern "C" BoxedGenerator* createGenerator(BoxedFunction* function, Box* arg1, B
 
 
 extern "C" BoxedGenerator::BoxedGenerator(BoxedFunction* function, Box* arg1, Box* arg2, Box* arg3, Box** args)
-    : Box(&generator_flavor, generator_cls), function(function), arg1(arg1), arg2(arg2), arg3(arg3), args(args),
+    : Box(&generator_flavor, generator_cls), function(function), arg1(arg1), arg2(arg2), arg3(arg3), args(nullptr),
       entryExited(false), returnValue(nullptr), exception(nullptr) {
 
-    function->generator = this; // HACK: this only alows one active generator
-
     giveAttr("__name__", boxString(function->f->source->getName()));
 
+    int numArgs = function->f->num_args;
+    if (numArgs > 3) {
+        numArgs -= 3;
+        this->args = new (numArgs) GCdArray();
+        memcpy(&this->args->elts[0], args, numArgs * sizeof(Box*));
+    }
+
     getcontext(&context);
     context.uc_link = 0;
     context.uc_stack.ss_sp = stack;

src/runtime/objmodel.cpp

@@ -41,6 +41,7 @@
 #include "runtime/capi.h"
 #include "runtime/float.h"
 #include "runtime/gc_runtime.h"
+#include "runtime/generator.h"
 #include "runtime/types.h"
 #include "runtime/util.h"
 
@@ -1717,7 +1718,6 @@ Box* callFunc(BoxedFunction* func, CallRewriteArgs* rewrite_args, ArgPassSpec ar
     int num_passed_args = argspec.totalPassed();
 
     BoxedClosure* closure = func->closure;
-    BoxedGenerator* generator = (BoxedGenerator*)func->generator;
 
     if (argspec.has_starargs || argspec.has_kwargs || f->takes_kwargs)
         rewrite_args = NULL;
@@ -1754,25 +1754,25 @@ Box* callFunc(BoxedFunction* func, CallRewriteArgs* rewrite_args, ArgPassSpec ar
 
     if (rewrite_args) {
         int closure_indicator = closure ? 1 : 0;
-        int generator_indicator = generator ? 1 : 0;
-        int argOffset = closure_indicator + generator_indicator;
+        int generator_indicator = func->isGenerator ? 1 : 0;
+        int arg_offset = closure_indicator + generator_indicator;
 
         if (num_passed_args >= 1)
-            rewrite_args->arg1 = rewrite_args->arg1.move(0 + argOffset);
+            rewrite_args->arg1 = rewrite_args->arg1.move(0 + arg_offset);
         if (num_passed_args >= 2)
-            rewrite_args->arg2 = rewrite_args->arg2.move(1 + argOffset);
+            rewrite_args->arg2 = rewrite_args->arg2.move(1 + arg_offset);
         if (num_passed_args >= 3)
-            rewrite_args->arg3 = rewrite_args->arg3.move(2 + argOffset);
+            rewrite_args->arg3 = rewrite_args->arg3.move(2 + arg_offset);
         if (num_passed_args >= 4)
-            rewrite_args->args = rewrite_args->args.move(3 + argOffset);
+            rewrite_args->args = rewrite_args->args.move(3 + arg_offset);
 
         // TODO this kind of embedded reference needs to be tracked by the GC somehow?
         // Or maybe it's ok, since we've guarded on the function object?
         if (closure)
             rewrite_args->rewriter->loadConst(0, (intptr_t)closure);
 
-        if (generator)
-            rewrite_args->rewriter->loadConst(0, (intptr_t)generator);
+        if (func->isGenerator)
+            rewrite_args->rewriter->loadConst(0, (intptr_t)0 /*generator*/);
 
         // We might have trouble if we have more output args than input args,
         // such as if we need more space to pass defaults.
@@ -1957,37 +1957,45 @@ Box* callFunc(BoxedFunction* func, CallRewriteArgs* rewrite_args, ArgPassSpec ar
         getArg(i, oarg1, oarg2, oarg3, oargs) = default_obj;
     }
 
+    // special handling for generators:
+    // the call to function containing a yield should just create a new generator object.
+    if (func->isGenerator) {
+        return createGenerator(func, oarg1, oarg2, oarg3, oargs);
+    }
 
+    return callCLFunc(f, rewrite_args, num_output_args, closure, NULL, oarg1, oarg2, oarg3, oargs);
+}
 
+Box* callCLFunc(CLFunction* f, CallRewriteArgs* rewrite_args, int num_output_args, BoxedClosure* closure,
+                BoxedGenerator* generator, Box* oarg1, Box* oarg2, Box* oarg3, Box** oargs) {
     CompiledFunction* chosen_cf = pickVersion(f, num_output_args, oarg1, oarg2, oarg3, oargs);
 
     assert(chosen_cf->is_interpreted == (chosen_cf->code == NULL));
     if (chosen_cf->is_interpreted) {
-        return interpretFunction(chosen_cf->func, num_output_args, func->closure, generator, oarg1, oarg2, oarg3,
-                                 oargs);
-    } else {
-        if (rewrite_args) {
-            rewrite_args->rewriter->addDependenceOn(chosen_cf->dependent_callsites);
-
-            RewriterVar var = rewrite_args->rewriter->call((void*)chosen_cf->call);
+        return interpretFunction(chosen_cf->func, num_output_args, closure, generator, oarg1, oarg2, oarg3, oargs);
+    }
 
-            rewrite_args->out_rtn = var;
-            rewrite_args->out_success = true;
-        }
+    if (rewrite_args) {
+        rewrite_args->rewriter->addDependenceOn(chosen_cf->dependent_callsites);
 
+        RewriterVar var = rewrite_args->rewriter->call((void*)chosen_cf->call);
 
-        if (closure && generator)
-            return chosen_cf->closure_generator_call(closure, generator, oarg1, oarg2, oarg3, oargs);
-        else if (closure)
-            return chosen_cf->closure_call(closure, oarg1, oarg2, oarg3, oargs);
-        else if (generator)
-            return chosen_cf->generator_call(generator, oarg1, oarg2, oarg3, oargs);
-        else
-            return chosen_cf->call(oarg1, oarg2, oarg3, oargs);
+        rewrite_args->out_rtn = var;
+        rewrite_args->out_success = true;
     }
+
+    if (closure && generator)
+        return chosen_cf->closure_generator_call(closure, generator, oarg1, oarg2, oarg3, oargs);
+    else if (closure)
+        return chosen_cf->closure_call(closure, oarg1, oarg2, oarg3, oargs);
+    else if (generator)
+        return chosen_cf->generator_call(generator, oarg1, oarg2, oarg3, oargs);
+    else
+        return chosen_cf->call(oarg1, oarg2, oarg3, oargs);
 }
 
 
+
 Box* runtimeCallInternal(Box* obj, CallRewriteArgs* rewrite_args, ArgPassSpec argspec, Box* arg1, Box* arg2, Box* arg3,
                          Box** args, const std::vector<const std::string*>* keyword_names) {
     int npassed_args = argspec.totalPassed();

src/runtime/objmodel.h

@@ -103,5 +103,8 @@ extern "C" void raiseNotIterableError(const char* typeName) __attribute__((__nor
 Box* typeCall(Box*, BoxedList*);
 Box* typeNew(Box*, Box*);
 bool isUserDefined(BoxedClass* cls);
+
+Box* callCLFunc(CLFunction* f, CallRewriteArgs* rewrite_args, int num_output_args, BoxedClosure* closure,
+                BoxedGenerator* generator, Box* oarg1, Box* oarg2, Box* oarg3, Box** oargs);
 }
 #endif

src/runtime/types.cpp

@@ -77,7 +77,7 @@ llvm::iterator_range<BoxIterator> Box::pyElements() {
 }
 
 extern "C" BoxedFunction::BoxedFunction(CLFunction* f)
-    : Box(&function_flavor, function_cls), f(f), closure(NULL), generator(nullptr), ndefaults(0), defaults(NULL) {
+    : Box(&function_flavor, function_cls), f(f), closure(NULL), isGenerator(false), ndefaults(0), defaults(NULL) {
     if (f->source) {
         assert(f->source->ast);
         // this->giveAttr("__name__", boxString(&f->source->ast->name));
@@ -91,8 +91,9 @@ extern "C" BoxedFunction::BoxedFunction(CLFunction* f)
 }
 
 extern "C" BoxedFunction::BoxedFunction(CLFunction* f, std::initializer_list<Box*> defaults, BoxedClosure* closure,
-                                        BoxedGenerator* generator)
-    : Box(&function_flavor, function_cls), f(f), closure(closure), generator(generator), ndefaults(0), defaults(NULL) {
+                                        bool isGenerator)
+    : Box(&function_flavor, function_cls), f(f), closure(closure), isGenerator(isGenerator), ndefaults(0),
+      defaults(NULL) {
     if (defaults.size()) {
         // make sure to initialize defaults first, since the GC behavior is triggered by ndefaults,
         // and a GC can happen within this constructor:
@@ -148,12 +149,12 @@ std::string BoxedModule::name() {
     }
 }
 
-extern "C" Box* boxCLFunction(CLFunction* f, BoxedClosure* closure, BoxedGenerator* generator,
+extern "C" Box* boxCLFunction(CLFunction* f, BoxedClosure* closure, bool isGenerator,
                               std::initializer_list<Box*> defaults) {
     if (closure)
         assert(closure->cls == closure_cls);
 
-    return new BoxedFunction(f, defaults, closure, generator);
+    return new BoxedFunction(f, defaults, closure, isGenerator);
 }
 
 extern "C" CLFunction* unboxCLFunction(Box* b) {
@@ -279,12 +280,30 @@ extern "C" void generatorGCHandler(GCVisitor* v, void* p) {
     boxGCHandler(v, p);
 
     BoxedGenerator* g = (BoxedGenerator*)p;
-    if (g->function)
+
+    if (g->function) {
         v->visit(g->function);
+        if (g->function->f) {
+            int num_args = g->function->f->num_args;
+
+            if (num_args >= 1)
+                v->visit(g->arg1);
+            if (num_args >= 2)
+                v->visit(g->arg2);
+            if (num_args >= 3)
+                v->visit(g->arg3);
+            if (num_args > 3)
+                v->visitPotentialRange(reinterpret_cast<void* const*>(&g->args->elts[0]),
+                                       reinterpret_cast<void* const*>(&g->args->elts[num_args - 3]));
+        }
+    }
     if (g->returnValue)
         v->visit(g->returnValue);
     if (g->exception)
         v->visit(g->exception);
+
+
+
     v->visitPotentialRange((void**)&g->context, ((void**)&g->context) + sizeof(g->context) / sizeof(void*));
     v->visitPotentialRange((void**)&g->returnContext,
                            ((void**)&g->returnContext) + sizeof(g->returnContext) / sizeof(void*));

src/runtime/types.h

@@ -92,7 +92,7 @@ Box* boxString(const std::string& s);
 extern "C" BoxedString* boxStrConstant(const char* chars);
 extern "C" void listAppendInternal(Box* self, Box* v);
 extern "C" void listAppendArrayInternal(Box* self, Box** v, int nelts);
-extern "C" Box* boxCLFunction(CLFunction* f, BoxedClosure* closure, BoxedGenerator* generator,
+extern "C" Box* boxCLFunction(CLFunction* f, BoxedClosure* closure, bool isGenerator,
                               std::initializer_list<Box*> defaults);
 extern "C" CLFunction* unboxCLFunction(Box* b);
 extern "C" Box* createUserClass(std::string* name, Box* base, Box* attr_dict);
@@ -280,14 +280,14 @@ public:
     HCAttrs attrs;
     CLFunction* f;
     BoxedClosure* closure;
-    BoxedGenerator* generator;
 
+    bool isGenerator;
     int ndefaults;
     GCdArray* defaults;
 
     BoxedFunction(CLFunction* f);
     BoxedFunction(CLFunction* f, std::initializer_list<Box*> defaults, BoxedClosure* closure = NULL,
-                  BoxedGenerator* generator = nullptr);
+                  bool isGenerator = false);
 };
 
 class BoxedModule : public Box {
@@ -332,7 +332,8 @@ public:
 
     HCAttrs attrs;
     BoxedFunction* function;
-    Box* arg1, *arg2, *arg3, **args;
+    Box* arg1, *arg2, *arg3;
+    GCdArray* args;
 
     bool entryExited;
     Box* returnValue;

test/tests/generators.py

@@ -13,9 +13,12 @@ def G2():
     yield 1
     yield 2
     yield 3
-g2 = G2()
-print list(g2)
-print list(g2)
+g2a = G2()
+g2b = G2()
+print g2b.next()
+print list(g2a)
+print list(g2b)
+print list(g2a)
 print list(G2())
 
 
@@ -67,3 +70,10 @@ def G6(a=[]):
 print list(G6())
 print list(G6())
 
+def G7(p):
+    a = p
+    b = 2
+    def G():
+        yield a+b
+    return G()
+print list(G7(1))