bjit: reenable local block var optimization

src/asm_writing/rewriter.cpp

@@ -1233,9 +1233,8 @@ void Rewriter::_call(RewriterVar* result, bool has_side_effects, void* func_addr
 std::vector<Location> Rewriter::getDecrefLocations() {
     std::vector<Location> decref_infos;
     for (RewriterVar& var : vars) {
-        if (var.locations.size() && var.needsDecref()) {
+        if (var.locations.size() && var.needsDecref(current_action_idx)) {
             bool found_location = false;
-
             for (Location l : var.locations) {
                 if (l.type == Location::Scratch) {
                     // convert to stack based location because later on we may not know the offset of the scratch area
@@ -1359,8 +1358,22 @@ void RewriterVar::refUsed() {
     rewriter->addAction([=]() { this->bumpUse(); }, { this }, ActionType::NORMAL);
 }
 
-bool RewriterVar::needsDecref() {
-    return reftype == RefType::OWNED && !this->refHandedOff();
+bool RewriterVar::needsDecref(int current_action_index) {
+    rewriter->assertPhaseEmitting();
+
+    if (reftype != RefType::OWNED)
+        return false;
+
+    // if nothing consumes this reference we need to create a decref entry
+    if (num_refs_consumed == 0)
+        return true;
+
+    // don't create decref entry if the currenty action hands off the ownership
+    int reference_handed_off_action_index = uses[last_refconsumed_numuses - 1];
+    if (reference_handed_off_action_index == current_action_index)
+        return false;
+
+    return true;
 }
 
 void RewriterVar::registerOwnedAttr(int byte_offset) {
@@ -1513,6 +1526,7 @@ void Rewriter::commit() {
             _incref(var, 1);
         }
 
+        current_action_idx = i;
         actions[i].action();
 
         if (failed) {
@@ -2215,6 +2229,7 @@ Rewriter::Rewriter(std::unique_ptr<ICSlotRewrite> rewrite, int num_args, const L
       return_location(this->rewrite->returnRegister()),
       failed(false),
       needs_invalidation_support(needs_invalidation_support),
+      current_action_idx(-1),
       added_changing_action(false),
       marked_inside_ic(false),
       done_guarding(false),

src/asm_writing/rewriter.h

@@ -254,7 +254,7 @@ private:
     bool isDoneUsing() { return next_use == uses.size(); }
     bool hasScratchAllocation() const { return scratch_allocation.second > 0; }
     void resetHasScratchAllocation() { scratch_allocation = std::make_pair(0, 0); }
-    bool needsDecref();
+    bool needsDecref(int current_action_index);
 
     // Indicates if this variable is an arg, and if so, what location the arg is from.
     bool is_arg;
@@ -432,6 +432,8 @@ protected:
              bool needs_invalidation_support = true);
 
     std::deque<RewriterAction> actions;
+    int current_action_idx; // in the emitting phase get's set to index of currently executed action
+
     template <typename F> RewriterAction* addAction(F&& action, llvm::ArrayRef<RewriterVar*> vars, ActionType type) {
         assertPhaseCollecting();
         for (RewriterVar* var : vars) {

src/codegen/ast_interpreter.cpp

@@ -471,13 +471,12 @@ void ASTInterpreter::doStore(AST_Name* node, STOLEN(Value) value) {
         bool closure = vst == ScopeInfo::VarScopeType::CLOSURE;
         if (jit) {
             bool is_live = true;
-            // TODO: turn this optimization back on.
-            // if (!closure)
-            // is_live = source_info->getLiveness()->isLiveAtEnd(name, current_block);
+            if (!closure)
+                is_live = source_info->getLiveness()->isLiveAtEnd(name, current_block);
             if (is_live)
                 jit->emitSetLocal(name, node->vreg, closure, value);
             else
-                jit->emitSetBlockLocal(name, value);
+                jit->emitSetBlockLocal(name, node->vreg, value);
         }
 
         if (closure) {
@@ -1781,6 +1780,10 @@ int ASTInterpreterJitInterface::getBoxedLocalsOffset() {
     return offsetof(ASTInterpreter, frame_info.boxedLocals);
 }
 
+int ASTInterpreterJitInterface::getCreatedClosureOffset() {
+    return offsetof(ASTInterpreter, created_closure);
+}
+
 int ASTInterpreterJitInterface::getCurrentBlockOffset() {
     return offsetof(ASTInterpreter, current_block);
 }

src/codegen/ast_interpreter.h

@@ -39,6 +39,7 @@ struct ASTInterpreterJitInterface {
     static constexpr uint64_t osr_dummy_value = -1;
 
     static int getBoxedLocalsOffset();
+    static int getCreatedClosureOffset();
     static int getCurrentBlockOffset();
     static int getCurrentInstOffset();
     static int getEdgeCountOffset();

src/codegen/baseline_jit.cpp

@@ -385,16 +385,15 @@ RewriterVar* JitFragmentWriter::emitGetBlockLocal(InternedString s, int vreg) {
     auto it = local_syms.find(s);
     if (it == local_syms.end()) {
         auto r = emitGetLocal(s, vreg);
-        // TODO: clear out the vreg?
-        // assert(r->reftype == RefType::OWNED);
-        // emitSetLocal(s, vreg, false, imm(nullptr));
-        // emitSetBlockLocal(s, r);
+        assert(r->reftype == RefType::OWNED);
+        emitSetBlockLocal(s, vreg, r);
         return r;
     }
     return it->second;
 }
 
 void JitFragmentWriter::emitKillTemporary(InternedString s, int vreg) {
+    if (!local_syms.count(s))
         emitSetLocal(s, vreg, false, imm(nullptr));
 }
 
@@ -549,8 +548,26 @@ std::vector<RewriterVar*> JitFragmentWriter::emitUnpackIntoArray(RewriterVar* v,
 }
 
 RewriterVar* JitFragmentWriter::emitYield(RewriterVar* v) {
-    auto rtn = call(false, (void*)ASTInterpreterJitInterface::yieldHelper, getInterp(), v)->setType(RefType::OWNED);
+    llvm::SmallVector<RewriterVar*, 16> local_args;
+    local_args.push_back(interp->getAttr(ASTInterpreterJitInterface::getCreatedClosureOffset()));
+    // we have to pass all owned references which are not stored in the vregs to yield() so that the GC can traverse it
+    for (auto&& sym : local_syms) {
+        if (sym.second == v)
+            continue;
+        if (sym.second->reftype == RefType::OWNED)
+            local_args.push_back(sym.second);
+    }
+    // erase duplicate entries
+    std::sort(local_args.begin(), local_args.end());
+    local_args.erase(std::unique(local_args.begin(), local_args.end()), local_args.end());
+
+    auto&& args = allocArgs(local_args, RewriterVar::SetattrType::REF_USED);
+    RewriterVar* generator = interp->getAttr(ASTInterpreterJitInterface::getGeneratorOffset());
+    auto rtn = call(false, (void*)yield, generator, v, args, imm(local_args.size()))->setType(RefType::OWNED);
     v->refConsumed();
+    for (auto&& var : local_args) {
+        var->refUsed();
+    }
     return rtn;
 }
 
@@ -643,10 +660,12 @@ void JitFragmentWriter::emitSetAttr(AST_expr* node, RewriterVar* obj, BoxedStrin
     attr->refConsumed(rtn.second);
 }
 
-void JitFragmentWriter::emitSetBlockLocal(InternedString s, STOLEN(RewriterVar*) v) {
+void JitFragmentWriter::emitSetBlockLocal(InternedString s, int vreg, STOLEN(RewriterVar*) v) {
     if (LOG_BJIT_ASSEMBLY)
         comment("BJIT: emitSetBlockLocal() start");
     RewriterVar* prev = local_syms[s];
+    if (!prev)
+        emitSetLocal(s, vreg, false, imm(nullptr)); // clear out the vreg
     local_syms[s] = v;
     if (LOG_BJIT_ASSEMBLY)
         comment("BJIT: emitSetBlockLocal() end");

src/codegen/baseline_jit.h

@@ -272,7 +272,7 @@ public:
     void emitRaise3(RewriterVar* arg0, RewriterVar* arg1, RewriterVar* arg2);
     void emitReturn(RewriterVar* v);
     void emitSetAttr(AST_expr* node, RewriterVar* obj, BoxedString* s, STOLEN(RewriterVar*) attr);
-    void emitSetBlockLocal(InternedString s, STOLEN(RewriterVar*) v);
+    void emitSetBlockLocal(InternedString s, int vreg, STOLEN(RewriterVar*) v);
     void emitSetCurrentInst(AST_stmt* node);
     void emitSetExcInfo(RewriterVar* type, RewriterVar* value, RewriterVar* traceback);
     void emitSetGlobal(BoxedString* s, STOLEN(RewriterVar*) v, bool are_globals_from_module);

test/tests/iter_del_time.py

-# fail-if: '-n' in EXTRA_JIT_ARGS or '-O' in EXTRA_JIT_ARGS
+# expected: fail
+# this only works in the interpreter and not in the bjit and llvm jit
 
 class C(object):
     def next(self):