Add support for generator arguments

kwargs are unsupported for now

src/analysis/function_analysis.cpp

@@ -31,25 +31,6 @@
 
 namespace pyston {
 
-class YieldVisitor : public NoopASTVisitor {
-public:
-    YieldVisitor() : containsYield(false) {}
-
-    virtual bool visit_yield(AST_Yield*) {
-        containsYield = true;
-        return true;
-    }
-
-    bool containsYield;
-};
-
-bool containsYield(AST* ast) {
-    YieldVisitor visitor;
-    ast->accept(&visitor);
-    return visitor.containsYield;
-}
-
-
 class LivenessBBVisitor : public NoopASTVisitor {
 public:
     typedef llvm::SmallSet<std::string, 4> StrSet;

src/analysis/function_analysis.h

@@ -77,8 +77,6 @@ public:
     bool isPotentiallyUndefinedAfter(const std::string& name, CFGBlock* block);
 };
 
-bool containsYield(AST* ast);
-
 LivenessAnalysis* computeLivenessInfo(CFG*);
 PhiAnalysis* computeRequiredPhis(const SourceInfo::ArgNames&, CFG*, LivenessAnalysis*, ScopeInfo* scope_Info);
 }

src/analysis/scoping_analysis.cpp

@@ -20,6 +20,24 @@
 
 namespace pyston {
 
+class YieldVisitor : public NoopASTVisitor {
+public:
+    YieldVisitor() : containsYield(false) {}
+
+    virtual bool visit_yield(AST_Yield*) {
+        containsYield = true;
+        return true;
+    }
+
+    bool containsYield;
+};
+
+bool containsYield(AST* ast) {
+    YieldVisitor visitor;
+    ast->accept(&visitor);
+    return visitor.containsYield;
+}
+
 static bool isCompilerCreatedName(const std::string& name) {
     return name[0] == '!' || name[0] == '#';
 }
@@ -388,6 +406,9 @@ ScopeInfo* ScopingAnalysis::analyzeSubtree(AST* node) {
 
     ScopeInfo* rtn = scopes[node];
     assert(rtn);
+
+    rtn->setTakesGenerator(containsYield(node));
+
     return rtn;
 }
 

src/analysis/scoping_analysis.h

@@ -24,12 +24,16 @@ class AST_Module;
 
 class ScopeInfo {
 public:
+    ScopeInfo() : isGeneratorValue(false) {}
     virtual ~ScopeInfo() {}
     virtual ScopeInfo* getParent() = 0;
 
     virtual bool createsClosure() = 0;
     virtual bool takesClosure() = 0;
 
+    virtual bool takesGenerator() { return isGeneratorValue; }
+    virtual void setTakesGenerator(bool b = true) { isGeneratorValue = b; }
+
     virtual bool refersToGlobal(const std::string& name) = 0;
     virtual bool refersToClosure(const std::string name) = 0;
     virtual bool saveInClosure(const std::string name) = 0;
@@ -38,6 +42,9 @@ public:
     // the metaclass constructor.
     // An error to call this on a non-classdef node.
     virtual const std::unordered_set<std::string>& getClassDefLocalNames() = 0;
+
+protected:
+    bool isGeneratorValue;
 };
 
 class ScopingAnalysis {
@@ -58,6 +65,8 @@ public:
 };
 
 ScopingAnalysis* runScopingAnalysis(AST_Module* m);
+
+bool containsYield(AST* ast);
 }
 
 #endif

src/codegen/codegen.h

@@ -70,7 +70,7 @@ struct GlobalState {
     llvm::Type* llvm_flavor_type, *llvm_flavor_type_ptr;
     llvm::Type* llvm_opaque_type;
     llvm::Type* llvm_str_type_ptr;
-    llvm::Type* llvm_clfunction_type_ptr, *llvm_closure_type_ptr, *llvm_func_type_ptr;
+    llvm::Type* llvm_clfunction_type_ptr, *llvm_closure_type_ptr, *llvm_generator_type_ptr;
     llvm::Type* llvm_module_type_ptr, *llvm_bool_type_ptr;
     llvm::Type* i1, *i8, *i8_ptr, *i32, *i64, *void_, *double_;
     llvm::Type* vector_ptr;

src/codegen/compvars.cpp

@@ -519,19 +519,31 @@ ConcreteCompilerVariable* UnknownType::nonzero(IREmitter& emitter, const OpInfo&
 }
 
 CompilerVariable* makeFunction(IREmitter& emitter, CLFunction* f, CompilerVariable* closure,
-                               const std::vector<ConcreteCompilerVariable*>& defaults) {
+                               CompilerVariable* generator, 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
 
     llvm::Value* closure_v;
-    ConcreteCompilerVariable* converted = NULL;
+    ConcreteCompilerVariable* convertedClosure = NULL;
     if (closure) {
-        converted = closure->makeConverted(emitter, closure->getConcreteType());
-        closure_v = converted->getValue();
+        convertedClosure = closure->makeConverted(emitter, closure->getConcreteType());
+        closure_v = convertedClosure->getValue();
     } else {
         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*));
@@ -548,11 +560,14 @@ CompilerVariable* makeFunction(IREmitter& emitter, CLFunction* f, CompilerVariab
     }
 
     llvm::Value* boxed = emitter.getBuilder()->CreateCall(
-        g.funcs.boxCLFunction, std::vector<llvm::Value*>{ embedConstantPtr(f, g.llvm_clfunction_type_ptr), closure_v,
-                                                          scratch, getConstantInt(defaults.size(), g.i64) });
-
-    if (converted)
-        converted->decvref(emitter);
+        g.funcs.boxCLFunction,
+        std::vector<llvm::Value*>{ embedConstantPtr(f, g.llvm_clfunction_type_ptr), closure_v, generator_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);
 }
 
@@ -1201,6 +1216,23 @@ public:
 } _CLOSURE;
 ConcreteCompilerType* CLOSURE = &_CLOSURE;
 
+class GeneratorType : public ConcreteCompilerType {
+public:
+    llvm::Type* llvmType() { return g.llvm_generator_type_ptr; }
+    std::string debugName() { return "generator"; }
+
+    virtual ConcreteCompilerType* getConcreteType() { return this; }
+    virtual ConcreteCompilerType* getBoxType() { return GENERATOR; }
+
+    virtual void drop(IREmitter& emitter, VAR* var) {
+        // pass
+    }
+    virtual void grab(IREmitter& emitter, VAR* var) {
+        // pass
+    }
+} _GENERATOR;
+ConcreteCompilerType* GENERATOR = &_GENERATOR;
+
 class StrConstantType : public ValuedCompilerType<std::string*> {
 public:
     std::string debugName() { return "str_constant"; }

src/codegen/compvars.h

@@ -29,7 +29,7 @@ class CompilerType;
 class IREmitter;
 
 extern ConcreteCompilerType* INT, *BOXED_INT, *FLOAT, *BOXED_FLOAT, *VOID, *UNKNOWN, *BOOL, *STR, *NONE, *LIST, *SLICE,
-    *MODULE, *DICT, *BOOL, *BOXED_BOOL, *BOXED_TUPLE, *SET, *CLOSURE;
+    *MODULE, *DICT, *BOOL, *BOXED_BOOL, *BOXED_TUPLE, *SET, *CLOSURE, *GENERATOR;
 extern CompilerType* UNDEF;
 
 class CompilerType {
@@ -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* makeFunction(IREmitter& emitter, CLFunction*, CompilerVariable* closure, CompilerVariable* generator,
                                const std::vector<ConcreteCompilerVariable*>& defaults);
 ConcreteCompilerVariable* undefVariable();
 CompilerVariable* makeTuple(const std::vector<CompilerVariable*>& elts);

src/codegen/irgen.cpp

@@ -513,7 +513,7 @@ static void emitBBs(IRGenState* irstate, const char* bb_type, GuardList& out_gua
             assert(strcmp("opt", bb_type) == 0);
 
             if (ENABLE_REOPT && effort < EffortLevel::MAXIMAL && source->ast != NULL
-                && source->ast->type != AST_TYPE::Module && !containsYield(source->ast)) {
+                && source->ast->type != AST_TYPE::Module) {
                 llvm::BasicBlock* preentry_bb
                     = llvm::BasicBlock::Create(g.context, "pre_entry", irstate->getLLVMFunction(),
                                                llvm_entry_blocks[source->cfg->getStartingBlock()]);
@@ -913,10 +913,14 @@ CompiledFunction* doCompile(SourceInfo* source, const OSREntryDescriptor* entry_
     int nargs = source->arg_names.totalParameters();
     ASSERT(nargs == spec->arg_types.size(), "%d %ld", nargs, spec->arg_types.size());
 
+
     std::vector<llvm::Type*> llvm_arg_types;
     if (source->scoping->getScopeInfoForNode(source->ast)->takesClosure())
         llvm_arg_types.push_back(g.llvm_closure_type_ptr);
 
+    if (source->scoping->getScopeInfoForNode(source->ast)->takesGenerator())
+        llvm_arg_types.push_back(g.llvm_generator_type_ptr);
+
     if (entry_descriptor == NULL) {
         for (int i = 0; i < nargs; i++) {
             if (i == 3) {

src/codegen/irgen/hooks.cpp

@@ -238,7 +238,7 @@ void compileAndRunModule(AST_Module* m, BoxedModule* bm) {
     }
 
     if (cf->is_interpreted)
-        interpretFunction(cf->func, 0, NULL, NULL, NULL, NULL, NULL);
+        interpretFunction(cf->func, 0, NULL, NULL, NULL, NULL, NULL, NULL);
     else
         ((void (*)())cf->code)();
 }

src/codegen/irgen/irgenerator.cpp

@@ -206,6 +206,7 @@ static std::vector<const std::string*>* getKeywordNameStorage(AST_Call* node) {
 
 static const std::string CREATED_CLOSURE_NAME = "!created_closure";
 static const std::string PASSED_CLOSURE_NAME = "!passed_closure";
+static const std::string PASSED_GENERATOR_NAME = "!passed_generator";
 
 class IRGeneratorImpl : public IRGenerator {
 private:
@@ -1019,29 +1020,21 @@ private:
 
     CompilerVariable* evalYield(AST_Yield* node, ExcInfo exc_info) {
         assert(state != PARTIAL);
-        /*
-                AST_Name genName;
-                genName.id = "!__GENERATOR__";
-                genName.ctx_type = AST_TYPE::Load;
-
-                CompilerVariable* generator = evalName(&genName, exc_info);
-                ConcreteCompilerVariable* convertedGenerator = generator->makeConverted(emitter,
-           generator->getBoxType());
-                generator->decvref(emitter);
-        */
+
+        CompilerVariable* generator = _getFake(PASSED_GENERATOR_NAME, false);
+        ConcreteCompilerVariable* convertedGenerator = generator->makeConverted(emitter, generator->getBoxType());
+
+
         CompilerVariable* value = node->value ? evalExpr(node->value, exc_info) : getNone();
         ConcreteCompilerVariable* convertedValue = value->makeConverted(emitter, value->getBoxType());
         value->decvref(emitter);
 
-        llvm::Value* rtn = emitter.createCall2(exc_info, g.funcs.yield, convertedValue->getValue(),
+        llvm::Value* rtn = emitter.createCall2(exc_info, g.funcs.yield, convertedGenerator->getValue(),
                                                convertedValue->getValue()).getInstruction();
-        // convertedGenerator->decvref(emitter);
+        convertedGenerator->decvref(emitter);
         convertedValue->decvref(emitter);
 
         return new ConcreteCompilerVariable(UNKNOWN, rtn, true);
-
-        // CompilerVariable* value = evalExpr(node->value, exc_info);
-        // return getNone();
     }
 
     ConcreteCompilerVariable* unboxVar(ConcreteCompilerType* t, llvm::Value* v, bool grabbed) {
@@ -1481,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, {});
+        CompilerVariable* func = makeFunction(emitter, cl, created_closure, 0, {});
 
         CompilerVariable* attr_dict = func->call(emitter, getEmptyOpInfo(exc_info), ArgPassSpec(0), {}, NULL);
 
@@ -1582,19 +1575,21 @@ private:
             assert(created_closure);
         }
 
-        CompilerVariable* func = makeFunction(emitter, cl, created_closure, defaults);
-
-        for (auto d : defaults) {
-            d->decvref(emitter);
-        }
+        CompilerVariable* func = makeFunction(emitter, cl, created_closure,
+                                              (ConcreteCompilerVariable*)scope_info->takesGenerator(), defaults);
 
-        if (containsYield(node)) {
+        if (scope_info->takesGenerator()) {
             ConcreteCompilerVariable* converted = func->makeConverted(emitter, func->getBoxType());
-            CLFunction* clFunc = boxRTFunction((void*)createGenerator, UNKNOWN, 1, 1, false, false);
-            func = makeFunction(emitter, clFunc, NULL, { converted });
+            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);
         }
 
+        for (auto d : defaults) {
+            d->decvref(emitter);
+        }
+
         return func;
     }
 
@@ -2039,7 +2034,8 @@ private:
     }
 
     bool allowableFakeEndingSymbol(const std::string& name) {
-        return startswith(name, "!is_defined") || name == PASSED_CLOSURE_NAME || name == CREATED_CLOSURE_NAME;
+        return startswith(name, "!is_defined") || name == PASSED_CLOSURE_NAME || name == CREATED_CLOSURE_NAME
+               || name == PASSED_GENERATOR_NAME;
     }
 
     void endBlock(State new_state) {
@@ -2148,6 +2144,7 @@ public:
         if (myblock->successors.size() == 0) {
             st->erase(CREATED_CLOSURE_NAME);
             st->erase(PASSED_CLOSURE_NAME);
+            st->erase(PASSED_GENERATOR_NAME);
             assert(st->size() == 0); // shouldn't have anything live if there are no successors!
             return EndingState(st, phi_st, curblock);
         } else if (myblock->successors.size() > 1) {
@@ -2179,6 +2176,8 @@ public:
                     ending_type = getPassedClosureType();
                 } else if (it->first == CREATED_CLOSURE_NAME) {
                     ending_type = getCreatedClosureType();
+                } else if (it->first == PASSED_GENERATOR_NAME) {
+                    ending_type = GENERATOR;
                 } else {
                     ending_type = types->getTypeAtBlockEnd(it->first, myblock);
                 }
@@ -2231,6 +2230,7 @@ public:
 
         llvm::Value* passed_closure = NULL;
         llvm::Function::arg_iterator AI = irstate->getLLVMFunction()->arg_begin();
+
         if (scope_info->takesClosure()) {
             passed_closure = AI;
             _setFake(PASSED_CLOSURE_NAME, new ConcreteCompilerVariable(getPassedClosureType(), AI, true));
@@ -2245,6 +2245,11 @@ public:
             _setFake(CREATED_CLOSURE_NAME, new ConcreteCompilerVariable(getCreatedClosureType(), new_closure, true));
         }
 
+        if (scope_info->takesGenerator()) {
+            _setFake(PASSED_GENERATOR_NAME, new ConcreteCompilerVariable(GENERATOR, AI, true));
+            ++AI;
+        }
+
 
         std::vector<llvm::Value*> python_parameters;
         for (int i = 0; i < arg_types.size(); i++) {

src/codegen/llvm_interpreter.cpp

@@ -257,7 +257,8 @@ const LineInfo* getLineInfoForInterpretedFrame(void* frame_ptr) {
     }
 }
 
-Box* interpretFunction(llvm::Function* f, int nargs, Box* closure, Box* arg1, Box* arg2, Box* arg3, Box** args) {
+Box* interpretFunction(llvm::Function* f, int nargs, Box* closure, Box* generator, Box* arg1, Box* arg2, Box* arg3,
+                       Box** args) {
     assert(f);
 
 #ifdef TIME_INTERPRETS
@@ -281,20 +282,24 @@ Box* interpretFunction(llvm::Function* f, int nargs, Box* closure, Box* arg1, Bo
 
     int arg_num = -1;
     int closure_indicator = closure ? 1 : 0;
+    int generator_indicator = generator ? 1 : 0;
+    int argOffset = closure_indicator + generator_indicator;
     for (llvm::Argument& arg : f->args()) {
         arg_num++;
 
         if (arg_num == 0 && closure)
             symbols.insert(std::make_pair(static_cast<llvm::Value*>(&arg), Val(closure)));
-        else if (arg_num == 0 + closure_indicator)
+        else if ((arg_num == 0 || (arg_num == 1 && closure)) && generator)
+            symbols.insert(std::make_pair(static_cast<llvm::Value*>(&arg), Val(generator)));
+        else if (arg_num == 0 + argOffset)
             symbols.insert(std::make_pair(static_cast<llvm::Value*>(&arg), Val(arg1)));
-        else if (arg_num == 1 + closure_indicator)
+        else if (arg_num == 1 + argOffset)
             symbols.insert(std::make_pair(static_cast<llvm::Value*>(&arg), Val(arg2)));
-        else if (arg_num == 2 + closure_indicator)
+        else if (arg_num == 2 + argOffset)
             symbols.insert(std::make_pair(static_cast<llvm::Value*>(&arg), Val(arg3)));
         else {
-            assert(arg_num == 3 + closure_indicator);
-            assert(f->getArgumentList().size() == 4 + closure_indicator);
+            assert(arg_num == 3 + argOffset);
+            assert(f->getArgumentList().size() == 4 + argOffset);
             assert(f->getArgumentList().back().getType() == g.llvm_value_type_ptr->getPointerTo());
             symbols.insert(std::make_pair(static_cast<llvm::Value*>(&arg), Val((int64_t)args)));
             // printf("loading %%4 with %p\n", (void*)args);

src/codegen/llvm_interpreter.h

@@ -25,7 +25,8 @@ class Box;
 class GCVisitor;
 class LineInfo;
 
-Box* interpretFunction(llvm::Function* f, int nargs, Box* closure, Box* arg1, Box* arg2, Box* arg3, Box** args);
+Box* interpretFunction(llvm::Function* f, int nargs, Box* closure, Box* generator, Box* arg1, Box* arg2, Box* arg3,
+                       Box** args);
 
 void gatherInterpreterRoots(GCVisitor* visitor);
 const LineInfo* getLineInfoForInterpretedFrame(void* frame_ptr);

src/codegen/runtime_hooks.cpp

@@ -142,8 +142,8 @@ void initGlobalFuncs(GlobalState& g) {
     g.llvm_closure_type_ptr = g.stdlib_module->getTypeByName("class.pyston::BoxedClosure")->getPointerTo();
     assert(g.llvm_closure_type_ptr);
 
-    g.llvm_func_type_ptr = g.stdlib_module->getTypeByName("class.pyston::BoxedFunction")->getPointerTo();
-    assert(g.llvm_func_type_ptr);
+    g.llvm_generator_type_ptr = g.stdlib_module->getTypeByName("class.pyston::BoxedGenerator")->getPointerTo();
+    assert(g.llvm_generator_type_ptr);
 
 #define GET(N) g.funcs.N = getFunc((void*)N, STRINGIFY(N))
 

src/codegen/runtime_hooks.h

@@ -32,7 +32,7 @@ struct GlobalFuncs {
 
     llvm::Value* boxInt, *unboxInt, *boxFloat, *unboxFloat, *boxStringPtr, *boxCLFunction, *unboxCLFunction,
         *boxInstanceMethod, *boxBool, *unboxBool, *createTuple, *createDict, *createList, *createSlice,
-        *createUserClass, *createClosure, *createGenerator;
+        *createUserClass, *createClosure, *createGenerator, *insideGenerator;
     llvm::Value* getattr, *setattr, *print, *nonzero, *binop, *compare, *augbinop, *unboxedLen, *getitem, *getclsattr,
         *yield, *getGlobal, *setitem, *delitem, *unaryop, *import, *repr, *isinstance;
 

src/core/types.h

@@ -166,6 +166,7 @@ struct FunctionSpecialization {
 };
 
 class BoxedClosure;
+class BoxedGenerator;
 struct CompiledFunction {
 private:
 public:
@@ -178,6 +179,8 @@ public:
     union {
         Box* (*call)(Box*, Box*, Box*, Box**);
         Box* (*closure_call)(BoxedClosure*, Box*, Box*, Box*, Box**);
+        Box* (*closure_generator_call)(BoxedClosure*, BoxedGenerator*, Box*, Box*, Box*, Box**);
+        Box* (*generator_call)(BoxedGenerator*, Box*, Box*, Box*, Box**);
         void* code;
     };
     llvm::Value* llvm_code; // the llvm callable.

src/runtime/builtin_modules/builtins.cpp

@@ -371,7 +371,8 @@ BoxedModule* builtins_module;
 
 // TODO looks like CPython and pypy put this into an "exceptions" module:
 BoxedClass* Exception, *AssertionError, *AttributeError, *TypeError, *NameError, *KeyError, *IndexError, *IOError,
-    *OSError, *ZeroDivisionError, *ValueError, *UnboundLocalError, *RuntimeError, *ImportError, *StopIteration;
+    *OSError, *ZeroDivisionError, *ValueError, *UnboundLocalError, *RuntimeError, *ImportError, *StopIteration,
+    *GeneratorExit;
 
 const ObjectFlavor exception_flavor(&boxGCHandler, NULL);
 Box* exceptionNew1(BoxedClass* cls) {
@@ -459,6 +460,7 @@ void setupBuiltins() {
     RuntimeError = makeBuiltinException(Exception, "RuntimeError");
     ImportError = makeBuiltinException(Exception, "ImportError");
     StopIteration = makeBuiltinException(Exception, "StopIteration");
+    GeneratorExit = makeBuiltinException(Exception, "GeneratorExit");
 
     repr_obj = new BoxedFunction(boxRTFunction((void*)repr, UNKNOWN, 1));
     builtins_module->giveAttr("repr", repr_obj);

src/runtime/generator.cpp

@@ -32,7 +32,6 @@
 
 namespace pyston {
 
-BoxedGenerator* g_gen; // HACK REMOVE ME!!!
 
 static void generatorEntry(BoxedGenerator* self) {
     assert(self->cls == generator_cls);
@@ -40,7 +39,9 @@ static void generatorEntry(BoxedGenerator* self) {
 
     try {
         // call body of the generator
-        runtimeCall(self->function, ArgPassSpec(0), 0, 0, 0, 0, 0);
+        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);
     } catch (Box* e) {
         // unhandled exception: propagate the exception to the caller
         self->exception = e;
@@ -85,19 +86,29 @@ Box* generatorThrow(Box* s, BoxedClass* e) {
     return generatorSend(self, None);
 }
 
+Box* generatorClose(Box* s) {
+    assert(s->cls == generator_cls);
+    BoxedGenerator* self = static_cast<BoxedGenerator*>(s);
+
+    // check if the generator already exited
+    if (self->entryExited)
+        return None;
+
+    return generatorThrow(self, GeneratorExit);
+}
+
 Box* generatorNext(Box* s) {
     return generatorSend(s, None);
 }
 
-extern "C" Box* yield(Box* obj, Box* value) {
-    obj = g_gen;
+extern "C" Box* yield(BoxedGenerator* obj, Box* value) {
     assert(obj->cls == generator_cls);
     BoxedGenerator* self = static_cast<BoxedGenerator*>(obj);
     self->returnValue = value;
 
     swapcontext(&self->context, &self->returnContext);
 
-    // if the generator receives a exception from the caller we have throw it
+    // if the generator receives a exception from the caller we have to throw it
     if (self->exception) {
         Box* exception = self->exception;
         self->exception = nullptr;
@@ -106,15 +117,19 @@ extern "C" Box* yield(Box* obj, Box* value) {
     return self->returnValue;
 }
 
-extern "C" BoxedGenerator* createGenerator(BoxedFunction* function) {
+
+extern "C" BoxedGenerator* createGenerator(BoxedFunction* function, Box* arg1, Box* arg2, Box* arg3, Box** args) {
     assert(function);
     assert(function->cls == function_cls);
-    return new BoxedGenerator(function);
+    return new BoxedGenerator(function, arg1, arg2, arg3, args);
 }
 
-extern "C" BoxedGenerator::BoxedGenerator(BoxedFunction* function)
-    : Box(&generator_flavor, generator_cls), function(function), entryExited(false), returnValue(nullptr),
-      exception(nullptr) {
+
+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),
+      entryExited(false), returnValue(nullptr), exception(nullptr) {
+
+    function->generator = this; // HACK: this only alows one active generator
 
     giveAttr("__name__", boxString(function->f->source->getName()));
 
@@ -123,8 +138,6 @@ extern "C" BoxedGenerator::BoxedGenerator(BoxedFunction* function)
     context.uc_stack.ss_sp = stack;
     context.uc_stack.ss_size = STACK_SIZE;
     makecontext(&context, (void (*)(void))generatorEntry, 1, this);
-
-    g_gen = this;
 }
 
 
@@ -134,6 +147,7 @@ void setupGenerator() {
     generator_cls->giveAttr("__iter__",
                             new BoxedFunction(boxRTFunction((void*)generatorIter, typeFromClass(generator_cls), 1)));
 
+    generator_cls->giveAttr("close", new BoxedFunction(boxRTFunction((void*)generatorClose, UNKNOWN, 1)));
     generator_cls->giveAttr("next", new BoxedFunction(boxRTFunction((void*)generatorNext, UNKNOWN, 1)));
     generator_cls->giveAttr("send", new BoxedFunction(boxRTFunction((void*)generatorSend, UNKNOWN, 2)));
     generator_cls->giveAttr("throw", new BoxedFunction(boxRTFunction((void*)generatorThrow, UNKNOWN, 2)));

src/runtime/generator.h

@@ -25,8 +25,8 @@ extern "C" const ObjectFlavor generator_flavor;
 
 void setupGenerator();
 
-extern "C" Box* yield(Box* obj, Box* value);
-extern "C" BoxedGenerator* createGenerator(BoxedFunction* function);
+extern "C" Box* yield(BoxedGenerator* obj, Box* value);
+extern "C" BoxedGenerator* createGenerator(BoxedFunction* function, Box* arg1, Box* arg2, Box* arg3, Box** args);
 }
 
 #endif

src/runtime/objmodel.cpp

@@ -818,6 +818,8 @@ Box* getattr_internal(Box* obj, const std::string& attr, bool check_cls, bool al
         raiseExcHelper(NameError, "free variable '%s' referenced before assignment in enclosing scope", attr.c_str());
     }
 
+    if (obj->cls == generator_cls) {
+    }
 
     if (allow_custom) {
         // Don't need to pass icentry args, since we special-case __getattribtue__ and __getattr__ to use
@@ -1715,6 +1717,7 @@ 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;
@@ -1751,21 +1754,26 @@ 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;
 
         if (num_passed_args >= 1)
-            rewrite_args->arg1 = rewrite_args->arg1.move(0 + closure_indicator);
+            rewrite_args->arg1 = rewrite_args->arg1.move(0 + argOffset);
         if (num_passed_args >= 2)
-            rewrite_args->arg2 = rewrite_args->arg2.move(1 + closure_indicator);
+            rewrite_args->arg2 = rewrite_args->arg2.move(1 + argOffset);
         if (num_passed_args >= 3)
-            rewrite_args->arg3 = rewrite_args->arg3.move(2 + closure_indicator);
+            rewrite_args->arg3 = rewrite_args->arg3.move(2 + argOffset);
         if (num_passed_args >= 4)
-            rewrite_args->args = rewrite_args->args.move(3 + closure_indicator);
+            rewrite_args->args = rewrite_args->args.move(3 + argOffset);
 
         // 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);
+
         // We might have trouble if we have more output args than input args,
         // such as if we need more space to pass defaults.
         if (num_output_args > 3 && num_output_args > argspec.totalPassed()) {
@@ -1955,7 +1963,8 @@ Box* callFunc(BoxedFunction* func, CallRewriteArgs* rewrite_args, ArgPassSpec ar
 
     assert(chosen_cf->is_interpreted == (chosen_cf->code == NULL));
     if (chosen_cf->is_interpreted) {
-        return interpretFunction(chosen_cf->func, num_output_args, func->closure, oarg1, oarg2, oarg3, oargs);
+        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);
@@ -1966,8 +1975,13 @@ Box* callFunc(BoxedFunction* func, CallRewriteArgs* rewrite_args, ArgPassSpec ar
             rewrite_args->out_success = true;
         }
 
-        if (closure)
+
+        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);
     }

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), ndefaults(0), defaults(NULL) {
+    : Box(&function_flavor, function_cls), f(f), closure(NULL), generator(nullptr), ndefaults(0), defaults(NULL) {
     if (f->source) {
         assert(f->source->ast);
         // this->giveAttr("__name__", boxString(&f->source->ast->name));
@@ -90,8 +90,9 @@ extern "C" BoxedFunction::BoxedFunction(CLFunction* f)
     assert(f->num_defaults == ndefaults);
 }
 
-extern "C" BoxedFunction::BoxedFunction(CLFunction* f, std::initializer_list<Box*> defaults, BoxedClosure* closure)
-    : Box(&function_flavor, function_cls), f(f), closure(closure), ndefaults(0), defaults(NULL) {
+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) {
     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:
@@ -147,11 +148,12 @@ std::string BoxedModule::name() {
     }
 }
 
-extern "C" Box* boxCLFunction(CLFunction* f, BoxedClosure* closure, std::initializer_list<Box*> defaults) {
+extern "C" Box* boxCLFunction(CLFunction* f, BoxedClosure* closure, BoxedGenerator* generator,
+                              std::initializer_list<Box*> defaults) {
     if (closure)
         assert(closure->cls == closure_cls);
 
-    return new BoxedFunction(f, defaults, closure);
+    return new BoxedFunction(f, defaults, closure, generator);
 }
 
 extern "C" CLFunction* unboxCLFunction(Box* b) {

src/runtime/types.h

@@ -92,7 +92,8 @@ 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, std::initializer_list<Box*> defaults);
+extern "C" Box* boxCLFunction(CLFunction* f, BoxedClosure* closure, BoxedGenerator* generator,
+                              std::initializer_list<Box*> defaults);
 extern "C" CLFunction* unboxCLFunction(Box* b);
 extern "C" Box* createUserClass(std::string* name, Box* base, Box* attr_dict);
 extern "C" double unboxFloat(Box* b);
@@ -279,12 +280,14 @@ public:
     HCAttrs attrs;
     CLFunction* f;
     BoxedClosure* closure;
+    BoxedGenerator* generator;
 
     int ndefaults;
     GCdArray* defaults;
 
     BoxedFunction(CLFunction* f);
-    BoxedFunction(CLFunction* f, std::initializer_list<Box*> defaults, BoxedClosure* closure = NULL);
+    BoxedFunction(CLFunction* f, std::initializer_list<Box*> defaults, BoxedClosure* closure = NULL,
+                  BoxedGenerator* generator = nullptr);
 };
 
 class BoxedModule : public Box {
@@ -329,6 +332,7 @@ public:
 
     HCAttrs attrs;
     BoxedFunction* function;
+    Box* arg1, *arg2, *arg3, **args;
 
     bool entryExited;
     Box* returnValue;
@@ -337,7 +341,7 @@ public:
     ucontext_t context, returnContext;
     char stack[STACK_SIZE];
 
-    BoxedGenerator(BoxedFunction* function);
+    BoxedGenerator(BoxedFunction* function, Box* arg1, Box* arg2, Box* arg3, Box** args);
 };
 
 extern "C" void boxGCHandler(GCVisitor* v, void* p);
@@ -347,7 +351,7 @@ Box* exceptionNew2(BoxedClass* cls, Box* message);
 
 extern BoxedClass* Exception, *AssertionError, *AttributeError, *TypeError, *NameError, *KeyError, *IndexError,
     *IOError, *OSError, *ZeroDivisionError, *ValueError, *UnboundLocalError, *RuntimeError, *ImportError,
-    *StopIteration;
+    *StopIteration, *GeneratorExit;
 
 // cls should be obj->cls.
 // Added as parameter because it should typically be available

test/tests/generators.py

-def G1():
-    i = 0
+def G1(i=0):
     while True:
         yield i
         i += i
@@ -21,8 +20,7 @@ print list(G2())
 
 
 
-def G3():
-    i = 0
+def G3(i=0):
     while True:
         got = (yield i**2)
         print "i=", i, "got=", got
@@ -36,13 +34,13 @@ for i in range(5):
 
 
 
-def G4():
+def G4(i=1):
     1/0
     while True:
         print "unreachable"
 
 try:
-    print list(G4())
+    print list(G4(0))
 except ZeroDivisionError:
     print "catched a ZeroDivisionError"
     
@@ -61,3 +59,11 @@ g5 = G5()
 for i in range(5):
     print g5.next()
 print g5.throw(ZeroDivisionError)
+
+def G6(a=[]):
+    for i in range(2):
+        a.append(i)
+        yield a
+print list(G6())
+print list(G6())
+