Fix closures and NAME lookups / sets

src/capi/typeobject.cpp

@@ -1957,7 +1957,7 @@ bool update_slot(BoxedClass* type, llvm::StringRef attr) noexcept {
     }
     if (ptrs[0] == NULL)
         return false; /* Not an attribute that affects any slots */
-    int r = update_subclasses(type, boxString(attr), update_slots_callback, (void*)ptrs);
+    int r = update_subclasses(type, autoDecref(boxString(attr)), update_slots_callback, (void*)ptrs);
 
     // TODO this is supposed to be a CAPI function!
     if (r)

src/codegen/ast_interpreter.cpp

@@ -169,6 +169,7 @@ public:
         }
 
         Py_DECREF(this->globals);
+        Py_XDECREF(this->created_closure);
     }
 
     llvm::DenseMap<InternedString, int>& getSymVRegMap() {
@@ -485,15 +486,14 @@ void ASTInterpreter::doStore(AST_Name* node, STOLEN(Value) value) {
                 jit->emitSetLocal(name, node->vreg, closure, value);
         }
 
-        assert(getSymVRegMap().count(name));
-        assert(getSymVRegMap()[name] == node->vreg);
-        Box* prev = vregs[node->vreg];
-        vregs[node->vreg] = value.o;
-        Py_XDECREF(prev);
-
         if (closure) {
-            assert(0 && "check refcounting");
-            created_closure->elts[scope_info->getClosureOffset(name)] = value.o;
+            ASTInterpreterJitInterface::setLocalClosureHelper(this, node->vreg, name, value.o);
+        } else {
+            assert(getSymVRegMap().count(name));
+            assert(getSymVRegMap()[name] == node->vreg);
+            Box* prev = vregs[node->vreg];
+            vregs[node->vreg] = value.o;
+            Py_XDECREF(prev);
         }
     }
 }
@@ -1578,7 +1578,6 @@ Value ASTInterpreter::visit_name(AST_Name* node) {
             return v;
         }
         case ScopeInfo::VarScopeType::DEREF: {
-            assert(0 && "check refcounting");
             return Value(ASTInterpreterJitInterface::derefHelper(this, node->id),
                          jit ? jit->emitDeref(node->id) : NULL);
         }
@@ -1610,7 +1609,6 @@ Value ASTInterpreter::visit_name(AST_Name* node) {
             RELEASE_ASSERT(0, "should be unreachable");
         }
         case ScopeInfo::VarScopeType::NAME: {
-            assert(0 && "check refcounting");
             Value v;
             if (jit)
                 v.var = jit->emitGetBoxedLocal(node->id.getBox());
@@ -1716,6 +1714,7 @@ Box* ASTInterpreterJitInterface::derefHelper(void* _interpreter, InternedString
     if (val == NULL) {
         raiseExcHelper(NameError, "free variable '%s' referenced before assignment in enclosing scope", s.c_str());
     }
+    Py_INCREF(val);
     return val;
 }
 
@@ -1740,10 +1739,9 @@ Box* ASTInterpreterJitInterface::landingpadHelper(void* _interpreter) {
     return rtn;
 }
 
-Box* ASTInterpreterJitInterface::setExcInfoHelper(void* _interpreter, Box* type, Box* value, Box* traceback) {
+void ASTInterpreterJitInterface::setExcInfoHelper(void* _interpreter, Box* type, Box* value, Box* traceback) {
     ASTInterpreter* interpreter = (ASTInterpreter*)_interpreter;
     interpreter->getFrameInfo()->exc = ExcInfo(type, value, traceback);
-    return None;
 }
 
 void ASTInterpreterJitInterface::setLocalClosureHelper(void* _interpreter, long vreg, InternedString id, Box* v) {
@@ -1751,14 +1749,15 @@ void ASTInterpreterJitInterface::setLocalClosureHelper(void* _interpreter, long
 
     assert(interpreter->getSymVRegMap().count(id));
     assert(interpreter->getSymVRegMap()[id] == vreg);
+    Box* prev = interpreter->vregs[vreg];
     interpreter->vregs[vreg] = v;
-    interpreter->created_closure->elts[interpreter->scope_info->getClosureOffset(id)] = v;
+    interpreter->created_closure->elts[interpreter->scope_info->getClosureOffset(id)] = incref(v);
+    Py_XDECREF(prev);
 }
 
-Box* ASTInterpreterJitInterface::uncacheExcInfoHelper(void* _interpreter) {
+void ASTInterpreterJitInterface::uncacheExcInfoHelper(void* _interpreter) {
     ASTInterpreter* interpreter = (ASTInterpreter*)_interpreter;
     interpreter->getFrameInfo()->exc = ExcInfo(NULL, NULL, NULL);
-    return None;
 }
 
 void ASTInterpreterJitInterface::raise0Helper(void* _interpreter) {

src/codegen/ast_interpreter.h

@@ -45,9 +45,9 @@ struct ASTInterpreterJitInterface {
     static Box* derefHelper(void* interp, InternedString s);
     static Box* doOSRHelper(void* interp, AST_Jump* node);
     static Box* landingpadHelper(void* interp);
-    static Box* setExcInfoHelper(void* interp, Box* type, Box* value, Box* traceback);
+    static void setExcInfoHelper(void* interp, Box* type, Box* value, Box* traceback);
     static void setLocalClosureHelper(void* interp, long vreg, InternedString id, Box* v);
-    static Box* uncacheExcInfoHelper(void* interp);
+    static void uncacheExcInfoHelper(void* interp);
     static void raise0Helper(void* interp) __attribute__((noreturn));
 };
 

src/codegen/baseline_jit.cpp

@@ -226,7 +226,7 @@ RewriterVar* JitFragmentWriter::emitCallattr(AST_expr* node, RewriterVar* obj, B
     if (keyword_names_var)
         call_args.push_back(keyword_names_var);
 
-    return call(false, (void*)callattrHelper, call_args);
+    return call(false, (void*)callattrHelper, call_args)->setType(RefType::OWNED);
 #endif
 }
 
@@ -239,29 +239,29 @@ RewriterVar* JitFragmentWriter::emitCreateDict(const llvm::ArrayRef<RewriterVar*
                                                const llvm::ArrayRef<RewriterVar*> values) {
     assert(keys.size() == values.size());
     if (keys.empty())
-        return call(false, (void*)createDict);
+        return call(false, (void*)createDict)->setType(RefType::OWNED);
     else
-        return call(false, (void*)createDictHelper, imm(keys.size()), allocArgs(keys), allocArgs(values));
+        return call(false, (void*)createDictHelper, imm(keys.size()), allocArgs(keys), allocArgs(values))->setType(RefType::OWNED);
 }
 
 RewriterVar* JitFragmentWriter::emitCreateList(const llvm::ArrayRef<RewriterVar*> values) {
     auto num = values.size();
     if (num == 0)
-        return call(false, (void*)createList);
+        return call(false, (void*)createList)->setType(RefType::OWNED);
     else
-        return call(false, (void*)createListHelper, imm(num), allocArgs(values));
+        return call(false, (void*)createListHelper, imm(num), allocArgs(values))->setType(RefType::OWNED);
 }
 
 RewriterVar* JitFragmentWriter::emitCreateSet(const llvm::ArrayRef<RewriterVar*> values) {
     auto num = values.size();
     if (num == 0)
-        return call(false, (void*)createSet);
+        return call(false, (void*)createSet)->setType(RefType::OWNED);
     else
-        return call(false, (void*)createSetHelper, imm(num), allocArgs(values));
+        return call(false, (void*)createSetHelper, imm(num), allocArgs(values))->setType(RefType::OWNED);
 }
 
 RewriterVar* JitFragmentWriter::emitCreateSlice(RewriterVar* start, RewriterVar* stop, RewriterVar* step) {
-    return call(false, (void*)createSlice, start, stop, step);
+    return call(false, (void*)createSlice, start, stop, step)->setType(RefType::OWNED);
 }
 
 RewriterVar* JitFragmentWriter::emitCreateTuple(const llvm::ArrayRef<RewriterVar*> values) {
@@ -270,13 +270,13 @@ RewriterVar* JitFragmentWriter::emitCreateTuple(const llvm::ArrayRef<RewriterVar
     if (num == 0) {
         r = imm(EmptyTuple)->setType(RefType::BORROWED);
     } else if (num == 1)
-        r = call(false, (void*)BoxedTuple::create1, values[0]);
+        r = call(false, (void*)BoxedTuple::create1, values[0])->setType(RefType::OWNED);
     else if (num == 2)
-        r = call(false, (void*)BoxedTuple::create2, values[0], values[1]);
+        r = call(false, (void*)BoxedTuple::create2, values[0], values[1])->setType(RefType::OWNED);
     else if (num == 3)
-        r = call(false, (void*)BoxedTuple::create3, values[0], values[1], values[2]);
+        r = call(false, (void*)BoxedTuple::create3, values[0], values[1], values[2])->setType(RefType::OWNED);
     else
-        r = call(false, (void*)createTupleHelper, imm(num), allocArgs(values));
+        r = call(false, (void*)createTupleHelper, imm(num), allocArgs(values))->setType(RefType::OWNED);
 
     return r;
 }
@@ -284,14 +284,15 @@ RewriterVar* JitFragmentWriter::emitCreateTuple(const llvm::ArrayRef<RewriterVar
 RewriterVar* JitFragmentWriter::emitDeref(InternedString s) {
     return call(false, (void*)ASTInterpreterJitInterface::derefHelper, getInterp(),
 #ifndef NDEBUG
-                imm(asUInt(s).first), imm(asUInt(s).second));
+                imm(asUInt(s).first), imm(asUInt(s).second))
 #else
-                imm(asUInt(s)));
+                imm(asUInt(s)))
 #endif
+        ->setType(RefType::OWNED);
 }
 
 RewriterVar* JitFragmentWriter::emitExceptionMatches(RewriterVar* v, RewriterVar* cls) {
-    return call(false, (void*)exceptionMatchesHelper, v, cls);
+    return call(false, (void*)exceptionMatchesHelper, v, cls)->setType(RefType::OWNED);
 }
 
 RewriterVar* JitFragmentWriter::emitGetAttr(RewriterVar* obj, BoxedString* s, AST_expr* node) {
@@ -308,7 +309,7 @@ RewriterVar* JitFragmentWriter::emitGetBlockLocal(InternedString s, int vreg) {
 RewriterVar* JitFragmentWriter::emitGetBoxedLocal(BoxedString* s) {
     RewriterVar* boxed_locals = emitGetBoxedLocals();
     RewriterVar* globals = getInterp()->getAttr(ASTInterpreterJitInterface::getGlobalsOffset());
-    return call(false, (void*)boxedLocalsGet, boxed_locals, imm(s), globals);
+    return call(false, (void*)boxedLocalsGet, boxed_locals, imm(s), globals)->setType(RefType::OWNED);
 }
 
 RewriterVar* JitFragmentWriter::emitGetBoxedLocals() {
@@ -429,7 +430,7 @@ RewriterVar* JitFragmentWriter::emitRuntimeCall(AST_expr* node, RewriterVar* obj
     if (keyword_names_var)
         call_args.push_back(keyword_names_var);
 
-    return call(false, (void*)runtimeCallHelper, call_args);
+    return call(false, (void*)runtimeCallHelper, call_args)->setType(RefType::OWNED);
 #endif
 }
 
@@ -496,7 +497,7 @@ void JitFragmentWriter::emitOSRPoint(AST_Jump* node) {
 void JitFragmentWriter::emitPrint(RewriterVar* dest, RewriterVar* var, bool nl) {
     if (LOG_BJIT_ASSEMBLY) comment("BJIT: emitPrint() start");
     if (!dest)
-        dest = call(false, (void*)getSysStdout);
+        dest = call(false, (void*)getSysStdout)->setType(RefType::BORROWED);
     if (!var)
         var = imm(0ul);
     call(false, (void*)printHelper, dest, var, imm(nl));
@@ -557,7 +558,6 @@ void JitFragmentWriter::emitSetLocal(InternedString s, int vreg, bool set_closur
     if (LOG_BJIT_ASSEMBLY) comment("BJIT: emitSetLocal() start");
     assert(vreg >= 0);
     if (set_closure) {
-        assert(0 && "check refcounting");
         call(false, (void*)ASTInterpreterJitInterface::setLocalClosureHelper, getInterp(), imm(vreg),
 #ifndef NDEBUG
              imm(asUInt(s).first), imm(asUInt(s).second),
@@ -565,6 +565,7 @@ void JitFragmentWriter::emitSetLocal(InternedString s, int vreg, bool set_closur
              imm(asUInt(s)),
 #endif
              v);
+        v->refConsumed();
     } else {
         RewriterVar* prev = vregs_array->getAttr(8 * vreg);
         vregs_array->setAttr(8 * vreg, v);

src/runtime/objmodel.cpp

@@ -333,6 +333,8 @@ extern "C" Box** unpackIntoArray(Box* obj, int64_t expected_size) {
         return &l->elts->elts[0];
     }
 
+    RELEASE_ASSERT(0, "I don't think this is safe since elts will die");
+
     std::vector<Box*> elts;
     for (auto e : obj->pyElements()) {
         elts.push_back(e);
@@ -2329,7 +2331,6 @@ Box* getattrInternalGeneric(Box* obj, BoxedString* attr, GetattrRewriteArgs* rew
                 return val;
             }
         } else {
-            RELEASE_ASSERT(0, "need to check the refcounting for this");
             // More complicated when obj is a type
             // We have to look up the attr in the entire
             // class hierarchy, and we also have to check if it is a descriptor,
@@ -2370,6 +2371,7 @@ Box* getattrInternalGeneric(Box* obj, BoxedString* attr, GetattrRewriteArgs* rew
                 if (!local_get) {
                     if (rewrite_args)
                         rewrite_args->setReturn(r_val, ReturnConvention::HAS_RETURN);
+                    Py_INCREF(val);
                     return val;
                 }
 

src/runtime/types.cpp

@@ -317,6 +317,9 @@ extern "C" BoxedFunctionBase::BoxedFunctionBase(FunctionMetadata* md, std::initi
       doc(NULL) {
     assert((!globals) == (!md->source || md->source->scoping->areGlobalsFromModule()));
 
+    Py_XINCREF(closure);
+    Py_XINCREF(globals);
+
     if (defaults.size()) {
         // HAX copy+modify the BoxedTuple constructor so that we can put NULLs into the tuple.
         // We are going to separately be careful to make sure that those NULLs don't escape

src/runtime/types.h

@@ -1094,7 +1094,7 @@ class BoxedStaticmethod : public Box {
 public:
     Box* sm_callable;
 
-    BoxedStaticmethod(Box* callable) : sm_callable(callable){};
+    BoxedStaticmethod(Box* callable) : sm_callable(callable) {}
 
     DEFAULT_CLASS_SIMPLE(staticmethod_cls, true);
 
@@ -1107,7 +1107,7 @@ class BoxedClassmethod : public Box {
 public:
     Box* cm_callable;
 
-    BoxedClassmethod(Box* callable) : cm_callable(callable){};
+    BoxedClassmethod(Box* callable) : cm_callable(callable) {}
 
     DEFAULT_CLASS_SIMPLE(classmethod_cls, true);
 
@@ -1125,16 +1125,12 @@ public:
 
     BoxedClosure(BoxedClosure* parent) : parent(parent) {}
 
+    // TODO: convert this to a var-object and use DEFAULT_CLASS_VAR_SIMPLE
     void* operator new(size_t size, size_t nelts) __attribute__((visibility("default"))) {
-        /*
-        BoxedClosure* rtn
-            = static_cast<BoxedClosure*>(gc_alloc(_PyObject_VAR_SIZE(closure_cls, nelts), gc::GCKind::PYTHON));
-            */
-        BoxedClosure* rtn
-            = static_cast<BoxedClosure*>(PyObject_MALLOC(sizeof(BoxedClosure) + nelts * sizeof(Box*)));
+        BoxedClosure* rtn = static_cast<BoxedClosure*>(_PyObject_GC_Malloc(sizeof(BoxedClosure) + nelts * sizeof(Box*)));
         rtn->nelts = nelts;
-        rtn->cls = closure_cls;
-        _Py_NewReference(rtn);
+        PyObject_INIT(rtn, closure_cls);                                                                               \
+        _PyObject_GC_TRACK(rtn);                                                                                   \
         memset((void*)rtn->elts, 0, sizeof(Box*) * nelts);
         return rtn;
     }