More refcounting fixes

src/asm_writing/rewriter.cpp

@@ -1753,6 +1753,7 @@ void Rewriter::_checkAndThrowCAPIException(RewriterVar* r, int64_t exc_val) {
     assertConsistent();
 }
 
+#ifndef NDEBUG
 void RewriterVar::addAssertLastActionAction(int num_left) {
     static int n = 0;
     n++;
@@ -1765,20 +1766,28 @@ void RewriterVar::addAssertLastActionAction(int num_left) {
         num_left++;
 
     rewriter->addAction([=]() {
+        if (this->skip_assert_last_action) {
+            this->skip_assert_last_action--;
+            return;
+        }
         if (this->next_use + num_left != this->uses.size()) {
             printf("n: %d\n", this_n);
             printf("this: %p\n", this);
             printf("uses: %p\n", &this->uses);
-            // for (auto idx : this->uses) {
-            // printf("Action: %p\n", rewriter->actions[idx].action.func);
-            // }
+            for (auto idx : this->uses) {
+                //  hax: reach inside SmallFunction to pull out the function pointer.
+                printf("Action %d: %p\n", idx, *(void**)&rewriter->actions[idx].action);
+            }
             printf("%ld\n", this->uses.size());
+            printf("Error: expected there to be %d uses left, but found %ld\n", num_left,
+                   this->uses.size() - this->next_use);
             raise(SIGTRAP);
         }
         RELEASE_ASSERT(this->next_use + num_left == this->uses.size(), "%d %d %ld\n", this->next_use, num_left,
                        this->uses.size());
     }, {}, ActionType::NORMAL);
 }
+#endif
 
 assembler::Indirect Rewriter::indirectFor(Location l) {
     assert(l.type == Location::Scratch || l.type == Location::Stack);

src/asm_writing/rewriter.h

@@ -209,6 +209,11 @@ enum class RefType {
 // This might make more sense as an inner class of Rewriter, but
 // you can't forward-declare that :/
 class RewriterVar {
+#ifndef NDEBUG
+private:
+    int skip_assert_last_action = 0;
+#endif
+
 public:
     typedef llvm::SmallVector<RewriterVar*, 8> SmallVector;
 
@@ -240,19 +245,33 @@ public:
             return setType(RefType::BORROWED);
 
         assert(this->reftype == RefType::BORROWED);
-        return this->incvref();
+
+        // Special case: the rewriter can "resurrect" constants, and produce an existing
+        // RewriterVar even though at the bjit / IC level it's a new variable.
+        if (this->vrefcount == 0) {
+#ifndef NDEBUG
+            skip_assert_last_action++;
+#endif
+            vrefcount++;
+            return this;
+        } else {
+            return this->incvref();
+        }
     }
 
     RewriterVar* incvref() {
         assert(reftype == RefType::OWNED || reftype == RefType::BORROWED);
-        assert(vrefcount > 0 || (reftype == RefType::BORROWED && vrefcount >= 0));
+        //assert(vrefcount > 0 || (reftype == RefType::BORROWED && vrefcount >= 0));
+        assert(vrefcount > 0);
         vrefcount++;
         return this;
     }
 
+#ifndef NDEBUG
     // Assert that there are exactly num_left actions get added after this call.
     // Meant for debugging.  It's a way to cross between the two rewriter phases.
     void addAssertLastActionAction(int num_left=0);
+#endif
 
     RewriterVar* decvref() {
         assert(reftype == RefType::OWNED || reftype == RefType::BORROWED);
@@ -726,6 +745,10 @@ public:
     void checkAndThrowCAPIException(RewriterVar* r, int64_t exc_val = 0);
 
     void abort();
+
+    // note: commitReturning decref all of the args variables, whereas commit() does not.
+    // This should probably be made more consistent, but functions that use commitReturning
+    // usually want this behavior but ones that use
     void commit();
     void commitReturning(RewriterVar* rtn);
     void commitReturningNonPython(RewriterVar* rtn);

src/codegen/ast_interpreter.cpp

@@ -1125,7 +1125,7 @@ Value ASTInterpreter::createFunction(AST* node, AST_arguments* args, const std::
         if (!passed_globals_var)
             passed_globals_var = jit->imm(0ul);
         rtn.var = jit->call(false, (void*)createFunctionFromMetadata, jit->imm(md), closure_var, passed_globals_var,
-                            defaults_var, jit->imm(args->defaults.size()));
+                            defaults_var, jit->imm(args->defaults.size()))->setType(RefType::OWNED);
     }
 
     rtn.o = createFunctionFromMetadata(md, closure, passed_globals, u.il);
@@ -1520,7 +1520,7 @@ Value ASTInterpreter::visit_num(AST_Num* node) {
     Py_INCREF(o);
     RewriterVar* v = NULL;
     if (jit) {
-        v = jit->imm(o)->setType(RefType::BORROWED);
+        v = jit->imm(o)->asBorrowed();
     }
     return Value(o, v);
 }
@@ -1683,8 +1683,11 @@ Value ASTInterpreter::visit_tuple(AST_Tuple* node) {
 
 Value ASTInterpreter::visit_attribute(AST_Attribute* node) {
     Value v = visit_expr(node->value);
-    return Value(pyston::getattr(v.o, node->attr.getBox()),
-                 jit ? jit->emitGetAttr(v, node->attr.getBox(), node) : NULL);
+    Value r(pyston::getattr(v.o, node->attr.getBox()), jit ? jit->emitGetAttr(v, node->attr.getBox(), node) : NULL);
+    Py_DECREF(v.o);
+    if (jit)
+        v.var->decvref();
+    return r;
 }
 }
 

src/codegen/baseline_jit.cpp

@@ -351,41 +351,42 @@ RewriterVar* JitFragmentWriter::emitGetLocal(InternedString s, int vreg) {
 }
 
 RewriterVar* JitFragmentWriter::emitGetPystonIter(RewriterVar* v) {
-    return call(false, (void*)getPystonIter, v);
+    return call(false, (void*)getPystonIter, v)->setType(RefType::OWNED);
 }
 
 RewriterVar* JitFragmentWriter::emitHasnext(RewriterVar* v) {
-    return call(false, (void*)hasnextHelper, v);
+    return call(false, (void*)hasnextHelper, v)->setType(RefType::OWNED);
 }
 
 RewriterVar* JitFragmentWriter::emitImportFrom(RewriterVar* module, BoxedString* name) {
-    return call(false, (void*)importFrom, module, imm(name));
+    return call(false, (void*)importFrom, module, imm(name))->setType(RefType::OWNED);
 }
 
 RewriterVar* JitFragmentWriter::emitImportName(int level, RewriterVar* from_imports, llvm::StringRef module_name) {
     return call(false, (void*)import, imm(level), from_imports, imm(const_cast<char*>(module_name.data())),
-                imm(module_name.size()));
+                imm(module_name.size()))->setType(RefType::OWNED);
 }
 
 RewriterVar* JitFragmentWriter::emitImportStar(RewriterVar* module) {
     RewriterVar* globals = getInterp()->getAttr(ASTInterpreterJitInterface::getGlobalsOffset());
-    return call(false, (void*)importStar, module, globals);
+    return call(false, (void*)importStar, module, globals)->setType(RefType::OWNED);
 }
 
 RewriterVar* JitFragmentWriter::emitLandingpad() {
-    return call(false, (void*)ASTInterpreterJitInterface::landingpadHelper, getInterp());
+    return call(false, (void*)ASTInterpreterJitInterface::landingpadHelper, getInterp())->setType(RefType::OWNED);
 }
 
 RewriterVar* JitFragmentWriter::emitNonzero(RewriterVar* v) {
+    // nonzeroHelper returns bool
     return call(false, (void*)nonzeroHelper, v);
 }
 
 RewriterVar* JitFragmentWriter::emitNotNonzero(RewriterVar* v) {
-    return call(false, (void*)notHelper, v);
+    return call(false, (void*)notHelper, v)->setType(RefType::OWNED);
 }
 
 RewriterVar* JitFragmentWriter::emitRepr(RewriterVar* v) {
-    return call(false, (void*)repr, v);
+    return call(false, (void*)repr, v)->setType(RefType::OWNED);
 }
 
 RewriterVar* JitFragmentWriter::emitRuntimeCall(AST_expr* node, RewriterVar* obj, ArgPassSpec argspec,
@@ -555,11 +556,11 @@ void JitFragmentWriter::emitSetExcInfo(RewriterVar* type, RewriterVar* value, Re
 }
 
 void JitFragmentWriter::emitSetGlobal(Box* global, BoxedString* s, RewriterVar* v) {
-    emitPPCall((void*)setGlobal, { imm(global), imm(s), v }, 2, 512)->setType(RefType::OWNED);
+    emitPPCall((void*)setGlobal, { imm(global), imm(s), v }, 2, 512);
 }
 
 void JitFragmentWriter::emitSetItem(RewriterVar* target, RewriterVar* slice, RewriterVar* value) {
-    emitPPCall((void*)setitem, { target, slice, value }, 2, 512)->setType(RefType::OWNED);
+    emitPPCall((void*)setitem, { target, slice, value }, 2, 512);
 }
 
 void JitFragmentWriter::emitSetItemName(BoxedString* s, RewriterVar* v) {

src/runtime/objmodel.cpp

@@ -2452,7 +2452,7 @@ bool dataDescriptorSetSpecialCases(Box* obj, Box* val, Box* descr, SetattrRewrit
 }
 
 template <Rewritable rewritable>
-void setattrGeneric(Box* obj, BoxedString* attr, Box* val, SetattrRewriteArgs* rewrite_args) {
+void setattrGeneric(Box* obj, BoxedString* attr, STOLEN(Box*) val, SetattrRewriteArgs* rewrite_args) {
     if (rewritable == NOT_REWRITABLE) {
         assert(!rewrite_args);
         rewrite_args = NULL;
@@ -2608,6 +2608,7 @@ extern "C" void setattr(Box* obj, BoxedString* attr, Box* attr_val) {
     RewriterVar* r_setattr;
 
     if (tp_setattro == instance_setattro) {
+        assert(0 && "check refcounting");
         if (rewriter.get()) {
             SetattrRewriteArgs rewrite_args(rewriter.get(), rewriter->getArg(0), rewriter->getArg(2));
             instanceSetattroInternal(obj, attr, attr_val, &rewrite_args);
@@ -2618,6 +2619,7 @@ extern "C" void setattr(Box* obj, BoxedString* attr, Box* attr_val) {
 
         return;
     } else if (tp_setattro != PyObject_GenericSetAttr) {
+        assert(0 && "check refcounting");
         static BoxedString* setattr_str = getStaticString("__setattr__");
         if (rewriter.get()) {
             GetattrRewriteArgs rewrite_args(rewriter.get(), rewriter->getArg(0)->getAttr(offsetof(Box, cls)),
@@ -2650,6 +2652,9 @@ extern "C" void setattr(Box* obj, BoxedString* attr, Box* attr_val) {
             SetattrRewriteArgs rewrite_args(rewriter.get(), rewriter->getArg(0), rewriter->getArg(2));
             setattrGeneric<REWRITABLE>(obj, attr, attr_val, &rewrite_args);
             if (rewrite_args.out_success) {
+                rewriter->getArg(0)->decvref();
+                rewriter->getArg(1)->decvref();
+                // arg 2 vref got consumed by setattrGeneric
                 rewriter->commit();
             }
         } else {
@@ -2658,6 +2663,7 @@ extern "C" void setattr(Box* obj, BoxedString* attr, Box* attr_val) {
         return;
     }
 
+    assert(0 && "check refcounting");
     if (rewriter.get()) {
         assert(setattr);
 
@@ -3550,10 +3556,7 @@ void rearrangeArgumentsInternal(ParamReceiveSpec paramspec, const ParamNames* pa
         rewrite_success = false; // default case
     }
 
-    // Super fast path:
-    if (argspec.num_keywords == 0 && !argspec.has_starargs && !paramspec.takes_varargs && !argspec.has_kwargs
-        && argspec.num_args == paramspec.num_args && !paramspec.takes_kwargs) {
-        rewrite_success = true;
+    auto propagate_args = [&]() {
         if (num_output_args >= 1)
             Py_INCREF(oarg1);
         if (num_output_args >= 2)
@@ -3582,6 +3585,13 @@ void rearrangeArgumentsInternal(ParamReceiveSpec paramspec, const ParamNames* pa
                 //}
             }
         }
+    };
+
+    // Super fast path:
+    if (argspec.num_keywords == 0 && !argspec.has_starargs && !paramspec.takes_varargs && !argspec.has_kwargs
+        && argspec.num_args == paramspec.num_args && !paramspec.takes_kwargs) {
+        rewrite_success = true;
+        propagate_args();
         return;
     }
 
@@ -3589,8 +3599,6 @@ void rearrangeArgumentsInternal(ParamReceiveSpec paramspec, const ParamNames* pa
     // django-admin test this covers something like 93% of all calls to callFunc.
     if (argspec.num_keywords == 0 && argspec.has_starargs == paramspec.takes_varargs && !argspec.has_kwargs
         && argspec.num_args == paramspec.num_args && (!paramspec.takes_kwargs || paramspec.kwargsIndex() < 3)) {
-        assert(0 && "check refcounting");
-        assert(!rewrite_args && "check refcounting");
 
         // TODO could also do this for empty varargs
         if (paramspec.takes_kwargs) {
@@ -3600,11 +3608,11 @@ void rearrangeArgumentsInternal(ParamReceiveSpec paramspec, const ParamNames* pa
             getArg(idx, oarg1, oarg2, oarg3, NULL) = NULL; // pass NULL for kwargs
             if (rewrite_args) {
                 if (idx == 0)
-                    rewrite_args->arg1 = rewrite_args->rewriter->loadConst(0);
+                    rewrite_args->arg1 = rewrite_args->rewriter->loadConst(0)->asBorrowed();
                 else if (idx == 1)
-                    rewrite_args->arg2 = rewrite_args->rewriter->loadConst(0);
+                    rewrite_args->arg2 = rewrite_args->rewriter->loadConst(0)->asBorrowed();
                 else if (idx == 2)
-                    rewrite_args->arg3 = rewrite_args->rewriter->loadConst(0);
+                    rewrite_args->arg3 = rewrite_args->rewriter->loadConst(0)->asBorrowed();
                 else
                     abort();
             }
@@ -3624,14 +3632,12 @@ void rearrangeArgumentsInternal(ParamReceiveSpec paramspec, const ParamNames* pa
                         ->addAttrGuard(offsetof(Box, cls), (intptr_t)tuple_cls);
                 }
                 rewrite_success = true;
-                if (num_output_args > 3)
-                    memcpy(oargs, args, sizeof(Box*) * (num_output_args - 3));
+                propagate_args();
                 return;
             }
         } else {
             rewrite_success = true;
-            if (num_output_args > 3)
-                memcpy(oargs, args, sizeof(Box*) * (num_output_args - 3));
+            propagate_args();
             return;
         }
     }
@@ -4286,10 +4292,11 @@ Box* callCLFunc(FunctionMetadata* md, CallRewriteArgs* rewrite_args, int num_out
                     arg_vec.push_back(rewrite_args->arg2);
 
                 if (S == CXX)
-                    rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretFunction, arg_vec);
+                    rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretFunction, arg_vec)
+                                                ->setType(RefType::OWNED);
                 else
-                    rewrite_args->out_rtn
-                        = rewrite_args->rewriter->call(true, (void*)astInterpretHelper2ArgsCapi, arg_vec);
+                    rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretHelper2ArgsCapi,
+                                                                         arg_vec)->setType(RefType::OWNED);
             } else {
                 // Hacky workaround: the rewriter can only pass arguments in registers, so use this helper function
                 // to unpack some of the additional arguments:
@@ -4305,9 +4312,11 @@ Box* callCLFunc(FunctionMetadata* md, CallRewriteArgs* rewrite_args, int num_out
                     arg_array->setAttr(24, rewrite_args->args);
 
                 if (S == CXX)
-                    rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretHelper, arg_vec);
+                    rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretHelper, arg_vec)
+                                                ->setType(RefType::OWNED);
                 else
-                    rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretHelperCapi, arg_vec);
+                    rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretHelperCapi, arg_vec)
+                                                ->setType(RefType::OWNED);
             }
 
             rewrite_args->out_success = true;

src/runtime/objmodel.h

@@ -107,7 +107,7 @@ extern "C" Box* createBoxedIterWrapperIfNeeded(Box* o);
 
 struct SetattrRewriteArgs;
 template <Rewritable rewritable>
-void setattrGeneric(Box* obj, BoxedString* attr, Box* val, SetattrRewriteArgs* rewrite_args);
+void setattrGeneric(Box* obj, BoxedString* attr, STOLEN(Box*) val, SetattrRewriteArgs* rewrite_args);
 
 struct BinopRewriteArgs;
 template <Rewritable rewritable>

src/runtime/traceback.cpp

@@ -104,6 +104,25 @@ void BoxedTraceback::here(LineInfo lineInfo, Box** tb) {
     *tb = new BoxedTraceback(std::move(lineInfo), *tb);
 }
 
+void BoxedTraceback::dealloc(Box* b) noexcept {
+    abort();
+}
+
+int BoxedTraceback::traverse(Box* self, visitproc visit, void *arg) noexcept {
+    BoxedTraceback* tb = static_cast<BoxedTraceback*>(self);
+
+    Py_VISIT(tb->tb_next);
+    Py_VISIT(tb->py_lines);
+    Py_VISIT(tb->line.file);
+    Py_VISIT(tb->line.func);
+
+    return 0;
+}
+
+int BoxedTraceback::clear(Box* self) noexcept {
+    abort();
+}
+
 void setupTraceback() {
     traceback_cls = BoxedClass::create(type_cls, object_cls, 0, 0, sizeof(BoxedTraceback), false, "traceback",
                                        (destructor)BoxedTraceback::dealloc, NULL, true,

src/runtime/traceback.h

@@ -31,7 +31,11 @@ public:
     LineInfo line;
     Box* py_lines;
 
-    BoxedTraceback(LineInfo line, Box* tb_next) : tb_next(tb_next), line(std::move(line)), py_lines(NULL) {}
+    BoxedTraceback(LineInfo line, Box* tb_next) : tb_next(tb_next), line(std::move(line)), py_lines(NULL) {
+        Py_INCREF(tb_next);
+        Py_INCREF(line.file);
+        Py_INCREF(line.func);
+    }
 
     DEFAULT_CLASS(traceback_cls);
 
@@ -42,15 +46,9 @@ public:
     // somewhat equivalent to PyTraceBack_Here
     static void here(LineInfo lineInfo, Box** tb);
 
-    static void dealloc(Box* b) noexcept {
-        Py_FatalError("unimplemented");
-    }
-    static int traverse(Box* self, visitproc visit, void *arg) noexcept {
-        Py_FatalError("unimplemented");
-    }
-    static int clear(Box* self) noexcept {
-        Py_FatalError("unimplemented");
-    }
+    static void dealloc(Box* b) noexcept;
+    static int traverse(Box* self, visitproc visit, void *arg) noexcept;
+    static int clear(Box* self) noexcept;
 };
 
 void printTraceback(Box* b);