llvm jit: fix vreg handling inside the LLVM jit

src/codegen/ast_interpreter.cpp

@@ -156,8 +156,9 @@ public:
         Py_XDECREF(frame_info.boxedLocals);
 
         int nvregs = getMD()->calculateNumVRegs();
-
-        for (int i = 0; i < nvregs; i++) {
+        int nvregs_user_visible = getMD()->calculateNumUserVisibleVRegs();
+        // skip the user visible ones because they will get decrefed in deinitFrame
+        for (int i = nvregs_user_visible; i < nvregs; i++) {
             Py_XDECREF(vregs[i]);
         }
 

src/codegen/codegen.cpp

@@ -87,6 +87,20 @@ int FunctionMetadata::calculateNumVRegs() {
     return cfg->sym_vreg_map.size();
 }
 
+int FunctionMetadata::calculateNumUserVisibleVRegs() {
+    SourceInfo* source_info = source.get();
+
+    CFG* cfg = source_info->cfg;
+    assert(cfg && "We don't calculate the CFG inside this function because it can raise an exception and its "
+                  "therefore not safe to call at every point");
+
+    if (!cfg->hasVregsAssigned()) {
+        ScopeInfo* scope_info = source->getScopeInfo();
+        cfg->assignVRegs(param_names, scope_info);
+    }
+    return cfg->sym_vreg_map_user_visible.size();
+}
+
 void FunctionMetadata::addVersion(CompiledFunction* compiled) {
     assert(compiled);
     assert((compiled->spec != NULL) + (compiled->entry_descriptor != NULL) == 1);

src/codegen/irgen/irgenerator.cpp

@@ -239,8 +239,7 @@ void IRGenState::setupFrameInfoVar(llvm::Value* passed_closure, llvm::Value* pas
         assert(al->isStaticAlloca());
 
         assert(!vregs);
-        getMD()->calculateNumVRegs();
-        int num_user_visible_vregs = getMD()->source->cfg->sym_vreg_map_user_visible.size();
+        int num_user_visible_vregs = getMD()->calculateNumUserVisibleVRegs();
         if (num_user_visible_vregs > 0) {
             auto* vregs_alloca
                 = builder.CreateAlloca(g.llvm_value_type_ptr, getConstantInt(num_user_visible_vregs), "vregs");
@@ -1827,7 +1826,7 @@ private:
         return rtn;
     }
 
-    template <typename GetLLVMValCB> void _setVRegIfUserVisible(InternedString name, GetLLVMValCB get_llvm_val_cb) {
+    template <typename GetLLVMValCB> void _setVRegIfUserVisible(InternedString name, GetLLVMValCB get_llvm_val_cb, CompilerVariable* prev) {
         auto cfg = irstate->getSourceInfo()->cfg;
         if (!cfg->hasVregsAssigned())
             irstate->getMD()->calculateNumVRegs();
@@ -1839,7 +1838,14 @@ private:
             // looks like this store don't have to be volatile because llvm knows that the vregs are visible thru the
             // FrameInfo which escapes.
             auto* gep = emitter.getBuilder()->CreateConstInBoundsGEP1_64(irstate->getVRegsVar(), vreg);
-            emitter.getBuilder()->CreateStore(get_llvm_val_cb(), gep);
+            if (prev) {
+                auto* old_value = emitter.getBuilder()->CreateLoad(gep);
+                emitter.setType(old_value, RefType::OWNED);
+            }
+
+            llvm::Value* new_val = get_llvm_val_cb();
+            auto* store = emitter.getBuilder()->CreateStore(new_val, gep);
+            emitter.refConsumed(new_val, store);
         }
     }
 
@@ -1875,8 +1881,8 @@ private:
         } else {
             // FAST or CLOSURE
 
-            CompilerVariable*& prev = symbol_table[name];
-            prev = val;
+            CompilerVariable* prev = symbol_table[name];
+            symbol_table[name] = val;
 
             // Clear out the is_defined name since it is now definitely defined:
             assert(!isIsDefinedName(name.s()));
@@ -1896,7 +1902,7 @@ private:
             }
 
             auto&& get_llvm_val = [&]() { return val->makeConverted(emitter, UNKNOWN)->getValue(); };
-            _setVRegIfUserVisible(name, get_llvm_val);
+            _setVRegIfUserVisible(name, get_llvm_val, prev);
         }
     }
 
@@ -2082,7 +2088,8 @@ private:
         // SyntaxError: can not delete variable 'x' referenced in nested scope
         assert(vst == ScopeInfo::VarScopeType::FAST);
 
-        _setVRegIfUserVisible(target->id, []() { return getNullPtr(g.llvm_value_type_ptr); });
+        CompilerVariable* prev = symbol_table.count(target->id) ? symbol_table[target->id] : NULL;
+        _setVRegIfUserVisible(target->id, []() { return getNullPtr(g.llvm_value_type_ptr); }, prev);
 
         if (symbol_table.count(target->id) == 0) {
             llvm::CallSite call = emitter.createCall(

src/codegen/unwinding.cpp

@@ -940,11 +940,7 @@ Box* FrameInfo::updateBoxedLocals() {
     // But we also need to loop through all the uninitialized variables that we have
     // access to and delete them from the locals dict
     if (frame_info->boxedLocals->cls == dict_cls) {
-        BoxedDict* boxed_locals = (BoxedDict*)frame_info->boxedLocals;
-        for (auto&& new_elem : d->d) {
-            assert(0 && "refcounting here is wrong if the key already exists");
-            boxed_locals->d[incref(new_elem.first)] = incref(new_elem.second);
-        }
+        PyDict_Update((BoxedDict*)frame_info->boxedLocals, d);
     } else {
         for (const auto& p : *d) {
             Box* varname = p.first;
@@ -953,7 +949,7 @@ Box* FrameInfo::updateBoxedLocals() {
         }
     }
 
-    return frame_info->boxedLocals;
+    return incref(frame_info->boxedLocals);
 }
 
 AST_stmt* PythonFrameIterator::getCurrentStatement() {

src/core/types.h

@@ -495,6 +495,7 @@ public:
                     ExceptionStyle exception_style = CXX);
 
     int calculateNumVRegs();
+    int calculateNumUserVisibleVRegs();
 
     // Helper function, meant for the C++ runtime, which allocates a FunctionMetadata object and calls addVersion
     // once to it.

src/runtime/frame.cpp

@@ -17,6 +17,7 @@
 
 #include "codegen/unwinding.h"
 #include "core/ast.h"
+#include "core/cfg.h"
 #include "runtime/types.h"
 
 namespace pyston {
@@ -163,12 +164,7 @@ public:
         BoxedFrame* f = static_cast<BoxedFrame*>(b);
 
         _PyObject_GC_UNTRACK(f);
-
-        Py_DECREF(f->_back);
-        Py_DECREF(f->_code);
-        Py_DECREF(f->_globals);
-        Py_DECREF(f->_locals);
-
+        clear(b);
         f->cls->tp_free(b);
     }
     static int traverse(Box* self, visitproc visit, void *arg) noexcept {
@@ -180,7 +176,13 @@ public:
         return 0;
     }
     static int clear(Box* self) noexcept {
-        Py_FatalError("unimplemented");
+        BoxedFrame* o = static_cast<BoxedFrame*>(self);
+        assert(o->hasExited());
+        Py_CLEAR(o->_back);
+        Py_CLEAR(o->_code);
+        Py_CLEAR(o->_globals);
+        Py_CLEAR(o->_locals);
+        return 0;
     }
 };
 
@@ -217,6 +219,12 @@ extern "C" void deinitFrame(FrameInfo* frame_info) {
         frame->handleFrameExit();
         Py_CLEAR(frame_info->frame_obj);
     }
+    if (frame_info->vregs) {
+        int num_user_visible_vregs = frame_info->md->calculateNumUserVisibleVRegs();
+        for (int i = 0; i < num_user_visible_vregs; i++) {
+            Py_XDECREF(frame_info->vregs[i]);
+        }
+    }
     Py_CLEAR(frame_info->boxedLocals);
 
     if (frame_info->exc.type) {