Support non-module-globals in the llvm tier

(motivated by namedtuple)

This involves two main changes:
- changing the calling convention to pass `globals` as an argument if needed
  (this only applies going into compiled code, it's already passed into the interpreter)
- changing the llvm irgenerator to use the new globals object

microbenchmarks/dict_globals_ubench.py

+d = dict(x=1, y=0)
+exec """
+def g():
+    global y
+    y += x
+""" in d
+
+def f():
+    g = d['g']
+    for i in xrange(1000000):
+        g()
+        g()
+        g()
+        g()
+        g()
+        g()
+        g()
+        g()
+        g()
+        g()
+        d['y'] += i
+f()
+print d['y']

microbenchmarks/namedtuple_ubench.py

+from collections import namedtuple
+
+NT = namedtuple("NT", "")
+
+def f():
+    C = NT
+    for i in xrange(1000000):
+        C()
+        C()
+        C()
+        C()
+        C()
+        C()
+        C()
+        C()
+        C()
+        C()
+f()
+

src/codegen/ast_interpreter.cpp

@@ -647,7 +647,7 @@ Value ASTInterpreter::visit_jump(AST_Jump* node) {
 }
 
 Box* ASTInterpreter::doOSR(AST_Jump* node) {
-    bool can_osr = ENABLE_OSR && !FORCE_INTERPRETER && source_info->scoping->areGlobalsFromModule();
+    bool can_osr = ENABLE_OSR && !FORCE_INTERPRETER;
     if (!can_osr)
         return NULL;
 
@@ -728,6 +728,9 @@ Box* ASTInterpreter::doOSR(AST_Jump* node) {
     if (created_closure)
         sorted_symbol_table[source_info->getInternedStrings().get(CREATED_CLOSURE_NAME)] = created_closure;
 
+    if (!source_info->scoping->areGlobalsFromModule())
+        sorted_symbol_table[source_info->getInternedStrings().get(PASSED_GLOBALS_NAME)] = globals;
+
     sorted_symbol_table[source_info->getInternedStrings().get(FRAME_INFO_PTR_NAME)] = (Box*)&frame_info;
 
     if (found_entry == nullptr) {
@@ -1721,7 +1724,7 @@ Box* astInterpretFunction(CLFunction* clfunc, int nargs, Box* closure, Box* gene
     SourceInfo* source_info = clfunc->source.get();
 
     assert((!globals) == source_info->scoping->areGlobalsFromModule());
-    bool can_reopt = ENABLE_REOPT && !FORCE_INTERPRETER && (globals == NULL);
+    bool can_reopt = ENABLE_REOPT && !FORCE_INTERPRETER;
 
     // If the cfg hasn't been computed yet, just conservatively say that it will be a big function.
     // It shouldn't matter, since the cfg should only be NULL if this is the first execution of this
@@ -1729,8 +1732,6 @@ Box* astInterpretFunction(CLFunction* clfunc, int nargs, Box* closure, Box* gene
     int num_blocks = source_info->cfg ? source_info->cfg->blocks.size() : 10000;
     int threshold = num_blocks <= 20 ? (REOPT_THRESHOLD_BASELINE / 3) : REOPT_THRESHOLD_BASELINE;
     if (unlikely(can_reopt && (FORCE_OPTIMIZE || !ENABLE_INTERPRETER || clfunc->times_interpreted > threshold))) {
-        assert(!globals);
-
         clfunc->times_interpreted = 0;
 
         EffortLevel new_effort = EffortLevel::MODERATE;
@@ -1756,15 +1757,24 @@ Box* astInterpretFunction(CLFunction* clfunc, int nargs, Box* closure, Box* gene
 
         UNAVOIDABLE_STAT_TIMER(t0, "us_timer_in_jitted_code");
         Box* r;
-        if (closure && generator)
-            r = optimized->closure_generator_call((BoxedClosure*)closure, (BoxedGenerator*)generator, arg1, arg2, arg3,
-                                                  args);
-        else if (closure)
-            r = optimized->closure_call((BoxedClosure*)closure, arg1, arg2, arg3, args);
-        else if (generator)
-            r = optimized->generator_call((BoxedGenerator*)generator, arg1, arg2, arg3, args);
-        else
+        Box* maybe_args[3];
+        int nmaybe_args = 0;
+        if (closure)
+            maybe_args[nmaybe_args++] = closure;
+        if (generator)
+            maybe_args[nmaybe_args++] = generator;
+        if (globals)
+            maybe_args[nmaybe_args++] = globals;
+        if (nmaybe_args == 0)
             r = optimized->call(arg1, arg2, arg3, args);
+        else if (nmaybe_args == 1)
+            r = optimized->call1(maybe_args[0], arg1, arg2, arg3, args);
+        else if (nmaybe_args == 2)
+            r = optimized->call2(maybe_args[0], maybe_args[1], arg1, arg2, arg3, args);
+        else {
+            assert(nmaybe_args == 3);
+            r = optimized->call3(maybe_args[0], maybe_args[1], maybe_args[2], arg1, arg2, arg3, args);
+        }
 
         if (optimized->exception_style == CXX)
             return r;

src/codegen/compvars.cpp

@@ -776,7 +776,7 @@ ConcreteCompilerVariable* UnknownType::hasnext(IREmitter& emitter, const OpInfo&
     return boolFromI1(emitter, rtn_val);
 }
 
-CompilerVariable* makeFunction(IREmitter& emitter, CLFunction* f, CompilerVariable* closure, Box* globals,
+CompilerVariable* makeFunction(IREmitter& emitter, CLFunction* f, CompilerVariable* closure, llvm::Value* globals,
                                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
@@ -805,14 +805,13 @@ CompilerVariable* makeFunction(IREmitter& emitter, CLFunction* f, CompilerVariab
         scratch = getNullPtr(g.llvm_value_type_ptr_ptr);
     }
 
-    assert(globals == NULL);
-    llvm::Value* globals_v = getNullPtr(g.llvm_value_type_ptr);
+    assert(globals);
 
     // We know this function call can't throw, so it's safe to use emitter.getBuilder()->CreateCall() rather than
     // emitter.createCall().
     llvm::Value* boxed = emitter.getBuilder()->CreateCall(
         g.funcs.boxCLFunction, std::vector<llvm::Value*>{ embedRelocatablePtr(f, g.llvm_clfunction_type_ptr), closure_v,
-                                                          globals_v, scratch, getConstantInt(defaults.size(), g.i64) });
+                                                          globals, scratch, getConstantInt(defaults.size(), g.i64) });
 
     if (convertedClosure)
         convertedClosure->decvref(emitter);

src/codegen/compvars.h

@@ -416,7 +416,7 @@ CompilerVariable* makeUnicode(Box*);
 #if 0
 CompilerVariable* makeUnicode(IREmitter& emitter, llvm::StringRef);
 #endif
-CompilerVariable* makeFunction(IREmitter& emitter, CLFunction*, CompilerVariable* closure, Box* globals,
+CompilerVariable* makeFunction(IREmitter& emitter, CLFunction*, CompilerVariable* closure, llvm::Value* globals,
                                const std::vector<ConcreteCompilerVariable*>& defaults);
 ConcreteCompilerVariable* undefVariable();
 CompilerVariable* makeTuple(const std::vector<CompilerVariable*>& elts);

src/codegen/irgen.cpp

@@ -442,6 +442,12 @@ static void emitBBs(IRGenState* irstate, TypeAnalysis* types, const OSREntryDesc
                 continue;
             }
 
+            if (p.first.s() == PASSED_GLOBALS_NAME) {
+                assert(!source->scoping->areGlobalsFromModule());
+                irstate->setGlobals(from_arg);
+                continue;
+            }
+
             ConcreteCompilerType* phi_type;
             phi_type = getTypeAtBlockStart(types, p.first, target_block);
 
@@ -607,6 +613,8 @@ static void emitBBs(IRGenState* irstate, TypeAnalysis* types, const OSREntryDesc
                 // (we manually added it during the calculation of osr_syms):
                 if (p.first.s() == FRAME_INFO_PTR_NAME)
                     continue;
+                if (p.first.s() == PASSED_GLOBALS_NAME)
+                    continue;
 
                 ConcreteCompilerType* analyzed_type = getTypeAtBlockStart(types, p.first, block);
 
@@ -1009,6 +1017,9 @@ CompiledFunction* doCompile(CLFunction* clfunc, SourceInfo* source, ParamNames*
         if (source->is_generator)
             llvm_arg_types.push_back(g.llvm_generator_type_ptr);
 
+        if (!source->scoping->areGlobalsFromModule())
+            llvm_arg_types.push_back(g.llvm_value_type_ptr);
+
         for (int i = 0; i < nargs; i++) {
             if (i == 3) {
                 llvm_arg_types.push_back(g.llvm_value_type_ptr->getPointerTo());

src/codegen/irgen/irgenerator.cpp

@@ -56,6 +56,7 @@ IRGenState::IRGenState(CLFunction* clfunc, CompiledFunction* cf, SourceInfo* sou
       scratch_space(NULL),
       frame_info(NULL),
       frame_info_arg(NULL),
+      globals(NULL),
       scratch_size(0) {
     assert(cf->func);
     assert(!cf->clfunc); // in this case don't need to pass in sourceinfo
@@ -244,6 +245,26 @@ ScopeInfo* IRGenState::getScopeInfoForNode(AST* node) {
     return source->scoping->getScopeInfoForNode(node);
 }
 
+void IRGenState::setGlobals(llvm::Value* globals) {
+    assert(!source_info->scoping->areGlobalsFromModule());
+    assert(!this->globals);
+    this->globals = globals;
+}
+
+llvm::Value* IRGenState::getGlobals() {
+    if (!globals) {
+        assert(source_info->scoping->areGlobalsFromModule());
+        this->globals = embedRelocatablePtr(source_info->parent_module, g.llvm_value_type_ptr);
+    }
+    return this->globals;
+}
+
+llvm::Value* IRGenState::getGlobalsIfCustom() {
+    if (source_info->scoping->areGlobalsFromModule())
+        return getNullPtr(g.llvm_value_type_ptr);
+    return getGlobals();
+}
+
 class IREmitterImpl : public IREmitter {
 private:
     IRGenState* irstate;
@@ -352,9 +373,6 @@ public:
         // even though we could allocate it in-line; maybe it's infrequently used enough that it's better
         // to not have it take up cache space.
 
-        RELEASE_ASSERT(irstate->getSourceInfo()->scoping->areGlobalsFromModule(),
-                       "jit doesn't support custom globals yet");
-
         builder->setEmitter(this);
         builder->SetInsertPoint(curblock);
     }
@@ -508,6 +526,7 @@ const std::string CREATED_CLOSURE_NAME = "#created_closure";
 const std::string PASSED_CLOSURE_NAME = "#passed_closure";
 const std::string PASSED_GENERATOR_NAME = "#passed_generator";
 const std::string FRAME_INFO_PTR_NAME = "#frame_info_ptr";
+const std::string PASSED_GLOBALS_NAME = "#passed_globals";
 
 bool isIsDefinedName(llvm::StringRef name) {
     return startswith(name, "!is_defined_");
@@ -741,7 +760,7 @@ private:
                 module->decvref(emitter);
 
                 llvm::Value* r = emitter.createCall2(unw_info, g.funcs.importStar, converted_module->getValue(),
-                                                     embedParentModulePtr());
+                                                     irstate->getGlobals());
                 CompilerVariable* v = new ConcreteCompilerVariable(UNKNOWN, r, true);
 
                 converted_module->decvref(emitter);
@@ -1072,14 +1091,14 @@ private:
             ICSetupInfo* pp = createGetGlobalIC(getOpInfoForNode(node, unw_info).getTypeRecorder());
 
             std::vector<llvm::Value*> llvm_args;
-            llvm_args.push_back(embedParentModulePtr());
+            llvm_args.push_back(irstate->getGlobals());
             llvm_args.push_back(embedRelocatablePtr(node->id.getBox(), g.llvm_boxedstring_type_ptr));
 
             llvm::Value* uncasted = emitter.createIC(pp, (void*)pyston::getGlobal, llvm_args, unw_info);
             llvm::Value* r = emitter.getBuilder()->CreateIntToPtr(uncasted, g.llvm_value_type_ptr);
             return new ConcreteCompilerVariable(UNKNOWN, r, true);
         } else {
-            llvm::Value* r = emitter.createCall2(unw_info, g.funcs.getGlobal, embedParentModulePtr(),
+            llvm::Value* r = emitter.createCall2(unw_info, g.funcs.getGlobal, irstate->getGlobals(),
                                                  embedRelocatablePtr(node->id.getBox(), g.llvm_boxedstring_type_ptr));
             return new ConcreteCompilerVariable(UNKNOWN, r, true);
         }
@@ -1147,7 +1166,7 @@ private:
         } else if (vst == ScopeInfo::VarScopeType::NAME) {
             llvm::Value* boxedLocals = irstate->getBoxedLocalsVar();
             llvm::Value* attr = embedRelocatablePtr(node->id.getBox(), g.llvm_boxedstring_type_ptr);
-            llvm::Value* module = embedParentModulePtr();
+            llvm::Value* module = irstate->getGlobals();
             llvm::Value* r = emitter.createCall3(unw_info, g.funcs.boxedLocalsGet, boxedLocals, attr, module);
             return new ConcreteCompilerVariable(UNKNOWN, r, true);
         } else {
@@ -1402,7 +1421,7 @@ private:
         // but since the classdef can't create its own closure, shouldn't need to explicitly
         // create that scope to pass the closure through.
         assert(irstate->getSourceInfo()->scoping->areGlobalsFromModule());
-        CompilerVariable* func = makeFunction(emitter, cl, created_closure, NULL, {});
+        CompilerVariable* func = makeFunction(emitter, cl, created_closure, irstate->getGlobalsIfCustom(), {});
 
         CompilerVariable* attr_dict = func->call(emitter, getEmptyOpInfo(unw_info), ArgPassSpec(0), {}, NULL);
 
@@ -1465,8 +1484,7 @@ private:
             assert(created_closure);
         }
 
-        assert(irstate->getSourceInfo()->scoping->areGlobalsFromModule());
-        CompilerVariable* func = makeFunction(emitter, cl, created_closure, NULL, defaults);
+        CompilerVariable* func = makeFunction(emitter, cl, created_closure, irstate->getGlobalsIfCustom(), defaults);
 
         for (auto d : defaults) {
             d->decvref(emitter);
@@ -1694,11 +1712,18 @@ private:
         assert(vst != ScopeInfo::VarScopeType::DEREF);
 
         if (vst == ScopeInfo::VarScopeType::GLOBAL) {
-            // TODO do something special here so that it knows to only emit a monomorphic inline cache?
-            auto parent_module = llvm::ConstantExpr::getPointerCast(embedParentModulePtr(), g.llvm_value_type_ptr);
-            ConcreteCompilerVariable* module = new ConcreteCompilerVariable(MODULE, parent_module, false);
-            module->setattr(emitter, getEmptyOpInfo(unw_info), name.getBox(), val);
-            module->decvref(emitter);
+            if (irstate->getSourceInfo()->scoping->areGlobalsFromModule()) {
+                auto parent_module = llvm::ConstantExpr::getPointerCast(embedParentModulePtr(), g.llvm_value_type_ptr);
+                ConcreteCompilerVariable* module = new ConcreteCompilerVariable(MODULE, parent_module, false);
+                module->setattr(emitter, getEmptyOpInfo(unw_info), name.getBox(), val);
+                module->decvref(emitter);
+            } else {
+                auto converted = val->makeConverted(emitter, val->getBoxType());
+                emitter.createCall3(unw_info, g.funcs.setGlobal, irstate->getGlobals(),
+                                    embedRelocatablePtr(name.getBox(), g.llvm_boxedstring_type_ptr),
+                                    converted->getValue());
+                converted->decvref(emitter);
+            }
         } else if (vst == ScopeInfo::VarScopeType::NAME) {
             // TODO inefficient
             llvm::Value* boxedLocals = irstate->getBoxedLocalsVar();
@@ -1825,7 +1850,7 @@ private:
         // We could patchpoint this or try to avoid the overhead, but this should only
         // happen when the assertion is actually thrown so I don't think it will be necessary.
         static BoxedString* AssertionError_str = internStringImmortal("AssertionError");
-        llvm_args.push_back(emitter.createCall2(unw_info, g.funcs.getGlobal, embedParentModulePtr(),
+        llvm_args.push_back(emitter.createCall2(unw_info, g.funcs.getGlobal, irstate->getGlobals(),
                                                 embedRelocatablePtr(AssertionError_str, g.llvm_boxedstring_type_ptr)));
 
         ConcreteCompilerVariable* converted_msg = NULL;
@@ -1906,7 +1931,7 @@ private:
         ScopeInfo::VarScopeType vst = scope_info->getScopeTypeOfName(target->id);
         if (vst == ScopeInfo::VarScopeType::GLOBAL) {
             // Can't use delattr since the errors are different:
-            emitter.createCall2(unw_info, g.funcs.delGlobal, embedParentModulePtr(),
+            emitter.createCall2(unw_info, g.funcs.delGlobal, irstate->getGlobals(),
                                 embedRelocatablePtr(target->id.getBox(), g.llvm_boxedstring_type_ptr));
             return;
         }
@@ -2700,6 +2725,11 @@ public:
             ++AI;
         }
 
+        if (!irstate->getSourceInfo()->scoping->areGlobalsFromModule()) {
+            irstate->setGlobals(AI);
+            ++AI;
+        }
+
         std::vector<llvm::Value*> python_parameters;
         for (int i = 0; i < arg_types.size(); i++) {
             assert(AI != irstate->getLLVMFunction()->arg_end());

src/codegen/irgen/irgenerator.h

@@ -49,6 +49,7 @@ extern const std::string CREATED_CLOSURE_NAME;
 extern const std::string PASSED_CLOSURE_NAME;
 extern const std::string PASSED_GENERATOR_NAME;
 extern const std::string FRAME_INFO_PTR_NAME;
+extern const std::string PASSED_GLOBALS_NAME;
 
 
 // Class that holds state of the current IR generation, that might not be local
@@ -70,6 +71,7 @@ private:
     llvm::Value* frame_info;
     llvm::Value* boxed_locals;
     llvm::Value* frame_info_arg;
+    llvm::Value* globals;
     int scratch_size;
 
 public:
@@ -107,6 +109,12 @@ public:
     ParamNames* getParamNames() { return param_names; }
 
     void setFrameInfoArgument(llvm::Value* v) { frame_info_arg = v; }
+
+    void setGlobals(llvm::Value* globals);
+    // Returns the custom globals, or the module if the globals come from the module.
+    llvm::Value* getGlobals();
+    // Returns the custom globals, or null if the globals come from the module.
+    llvm::Value* getGlobalsIfCustom();
 };
 
 // turns CFGBlocks into LLVM IR

src/codegen/runtime_hooks.cpp

@@ -209,6 +209,7 @@ void initGlobalFuncs(GlobalState& g) {
     GET(delitem);
     GET(getGlobal);
     GET(delGlobal);
+    GET(setGlobal);
     GET(binop);
     GET(compare);
     GET(augbinop);

src/codegen/runtime_hooks.h

@@ -37,7 +37,8 @@ struct GlobalFuncs {
         *createGenerator, *createSet;
     llvm::Value* getattr, *getattr_capi, *setattr, *delattr, *delitem, *delGlobal, *nonzero, *binop, *compare,
         *augbinop, *unboxedLen, *getitem, *getitem_capi, *getclsattr, *getGlobal, *setitem, *unaryop, *import,
-        *importFrom, *importStar, *repr, *str, *strOrUnicode, *exceptionMatches, *yield, *getiterHelper, *hasnext;
+        *importFrom, *importStar, *repr, *str, *strOrUnicode, *exceptionMatches, *yield, *getiterHelper, *hasnext,
+        *setGlobal;
 
     llvm::Value* unpackIntoArray, *raiseAttributeError, *raiseAttributeErrorStr, *raiseAttributeErrorCapi,
         *raiseAttributeErrorStrCapi, *raiseNotIterableError, *raiseIndexErrorStr, *raiseIndexErrorStrCapi,

src/core/types.h

@@ -267,6 +267,9 @@ public:
         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**);
+        Box* (*call1)(Box*, Box*, Box*, Box*, Box**);
+        Box* (*call2)(Box*, Box*, Box*, Box*, Box*, Box**);
+        Box* (*call3)(Box*, Box*, Box*, Box*, Box*, Box*, Box**);
         void* code;
         uintptr_t code_start;
     };

src/runtime/inline/link_forcer.cpp

@@ -87,6 +87,7 @@ void force() {
     FORCE(getclsattr);
     FORCE(getGlobal);
     FORCE(delGlobal);
+    FORCE(setGlobal);
     FORCE(setitem);
     FORCE(delitem);
     FORCE(unaryop);

src/runtime/objmodel.cpp

@@ -3669,32 +3669,46 @@ Box* callFunc(BoxedFunctionBase* func, CallRewriteArgs* rewrite_args, ArgPassSpe
 }
 
 template <ExceptionStyle S>
-static Box* callChosenCF(CompiledFunction* chosen_cf, BoxedClosure* closure, BoxedGenerator* generator, Box* oarg1,
-                         Box* oarg2, Box* oarg3, Box** oargs) noexcept(S == CAPI) {
+static Box* callChosenCF(CompiledFunction* chosen_cf, BoxedClosure* closure, BoxedGenerator* generator, Box* globals,
+                         Box* oarg1, Box* oarg2, Box* oarg3, Box** oargs) noexcept(S == CAPI) {
     if (S != chosen_cf->exception_style) {
         if (S == CAPI) {
             try {
-                return callChosenCF<CXX>(chosen_cf, closure, generator, oarg1, oarg2, oarg3, oargs);
+                return callChosenCF<CXX>(chosen_cf, closure, generator, globals, oarg1, oarg2, oarg3, oargs);
             } catch (ExcInfo e) {
                 setCAPIException(e);
                 return NULL;
             }
         } else {
-            Box* r = callChosenCF<CAPI>(chosen_cf, closure, generator, oarg1, oarg2, oarg3, oargs);
+            Box* r = callChosenCF<CAPI>(chosen_cf, closure, generator, globals, oarg1, oarg2, oarg3, oargs);
             if (!r)
                 throwCAPIException();
             return r;
         }
     }
 
-    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
+    assert((globals == NULL)
+           == (!chosen_cf->clfunc->source || chosen_cf->clfunc->source->scoping->areGlobalsFromModule()));
+
+    Box* maybe_args[3];
+    int nmaybe_args = 0;
+    if (closure)
+        maybe_args[nmaybe_args++] = closure;
+    if (generator)
+        maybe_args[nmaybe_args++] = generator;
+    if (globals)
+        maybe_args[nmaybe_args++] = globals;
+
+    if (nmaybe_args == 0)
         return chosen_cf->call(oarg1, oarg2, oarg3, oargs);
+    else if (nmaybe_args == 1)
+        return chosen_cf->call1(maybe_args[0], oarg1, oarg2, oarg3, oargs);
+    else if (nmaybe_args == 2)
+        return chosen_cf->call2(maybe_args[0], maybe_args[1], oarg1, oarg2, oarg3, oargs);
+    else {
+        assert(nmaybe_args == 3);
+        return chosen_cf->call3(maybe_args[0], maybe_args[1], maybe_args[2], oarg1, oarg2, oarg3, oargs);
+    }
 }
 
 // This function exists for the rewriter: astInterpretFunction takes 9 args, but the rewriter
@@ -3786,8 +3800,6 @@ Box* callCLFunc(CLFunction* f, CallRewriteArgs* rewrite_args, int num_output_arg
         }
     }
 
-    ASSERT(!globals, "need to update the calling conventions if we want to pass globals");
-
     if (rewrite_args) {
         rewrite_args->rewriter->addDependenceOn(chosen_cf->dependent_callsites);
 
@@ -3803,6 +3815,8 @@ Box* callCLFunc(CLFunction* f, CallRewriteArgs* rewrite_args, int num_output_arg
 
         if (closure)
             arg_vec.push_back(rewrite_args->rewriter->loadConst((intptr_t)closure, Location::forArg(0)));
+        if (globals)
+            arg_vec.push_back(rewrite_args->rewriter->loadConst((intptr_t)globals, Location::forArg(0)));
         if (num_output_args >= 1)
             arg_vec.push_back(rewrite_args->arg1);
         if (num_output_args >= 2)
@@ -3840,10 +3854,10 @@ Box* callCLFunc(CLFunction* f, CallRewriteArgs* rewrite_args, int num_output_arg
     // code and calls that target to builtins.
     if (f->source) {
         UNAVOIDABLE_STAT_TIMER(t0, "us_timer_in_jitted_code");
-        r = callChosenCF<S>(chosen_cf, closure, generator, oarg1, oarg2, oarg3, oargs);
+        r = callChosenCF<S>(chosen_cf, closure, generator, globals, oarg1, oarg2, oarg3, oargs);
     } else {
         UNAVOIDABLE_STAT_TIMER(t0, "us_timer_in_builtins");
-        r = callChosenCF<S>(chosen_cf, closure, generator, oarg1, oarg2, oarg3, oargs);
+        r = callChosenCF<S>(chosen_cf, closure, generator, globals, oarg1, oarg2, oarg3, oargs);
     }
 
     if (!r) {
@@ -5665,7 +5679,7 @@ extern "C" Box* getGlobal(Box* globals, BoxedString* name) {
                 }
             }
         } else {
-            assert(globals->cls == dict_cls);
+            ASSERT(globals->cls == dict_cls, "%s", globals->cls->tp_name);
             BoxedDict* d = static_cast<BoxedDict*>(globals);
 
             rewriter.reset(NULL);
@@ -5725,7 +5739,7 @@ Box* getFromGlobals(Box* globals, BoxedString* name) {
     }
 }
 
-void setGlobal(Box* globals, BoxedString* name, Box* value) {
+extern "C" void setGlobal(Box* globals, BoxedString* name, Box* value) {
     if (globals->cls == attrwrapper_cls) {
         globals = unwrapAttrWrapper(globals);
         RELEASE_ASSERT(globals->cls == module_cls, "%s", globals->cls->tp_name);

src/runtime/objmodel.h

@@ -211,7 +211,7 @@ extern "C" Box* getGlobal(Box* globals, BoxedString* name);
 // Checks for the name just in the passed globals object, and returns NULL if it is not found.
 // This includes if the globals object defined a custom __getattr__ method that threw an AttributeError.
 Box* getFromGlobals(Box* globals, BoxedString* name);
-void setGlobal(Box* globals, BoxedString* name, Box* value);
+extern "C" void setGlobal(Box* globals, BoxedString* name, Box* value);
 extern "C" void delGlobal(Box* globals, BoxedString* name);
 
 extern "C" void boxedLocalsSet(Box* boxedLocals, BoxedString* attr, Box* val);

test/tests/exec_in_test.py

@@ -199,3 +199,22 @@ print l
 
 exec s
 print __doc__
+
+
+# Create a function that needs all three extra arguments:
+# is a generator, takes a closure, and takes custom globals
+s = """
+def f(x):
+    def g(a, b, c, d, e):
+        for i in xrange(start, x):
+            print a, b, c, d, e
+            yield i
+    return g
+"""
+
+g = {'start':2}
+l = {}
+exec s in g, l
+
+for i in xrange(5):
+    print list(l['f'](5)(1, 2, 3, 4, 5))

test/tests/generator_custom_globals.py

+s = """
+def g():
+    yield x
+    yield y
+"""
+
+g = {'x': 1, 'y': 4}
+l = {}
+
+exec s in g, l
+
+print list(l['g']())