Rewriter setattr checking

A kind-of hacky way of identifying places that do unsafe (for refcounting)
setattrs.  When calling RewriterVar::setattr() with an owned reference,
you need to either promise to call refUsed() or refConsumed() afterwards.

src/asm_writing/rewriter.cpp

@@ -681,9 +681,15 @@ void Rewriter::_toBool(RewriterVar* result, RewriterVar* var, Location dest) {
     assertConsistent();
 }
 
-void RewriterVar::setAttr(int offset, RewriterVar* val) {
+void RewriterVar::setAttr(int offset, RewriterVar* val, SetattrType type) {
     STAT_TIMER(t0, "us_timer_rewriter", 10);
 
+    // Check that the caller promises to handle lifetimes appropriately.
+    // We're only interested in OWNED references, since we are trying to
+    // prevent store-in-array-and-pass situations where the refcounter will
+    // decref between the store and the pass.
+    if (val->reftype == RefType::OWNED)
+        assert(type != SetattrType::UNKNOWN);
     rewriter->addAction([=]() { rewriter->_setAttr(this, offset, val); }, { this, val }, ActionType::MUTATION);
 }
 

src/asm_writing/rewriter.h

@@ -149,7 +149,12 @@ public:
     // getAttrFloat casts to double (maybe I should make that separate?)
     RewriterVar* getAttrFloat(int offset, Location loc = Location::any());
     RewriterVar* getAttrDouble(int offset, Location loc = Location::any());
-    void setAttr(int offset, RewriterVar* other);
+    enum class SetattrType {
+        UNKNOWN,
+        HANDED_OFF,
+        REFUSED,
+    };
+    void setAttr(int offset, RewriterVar* other, SetattrType type = SetattrType::UNKNOWN);
     RewriterVar* cmp(AST_TYPE::AST_TYPE cmp_type, RewriterVar* other, Location loc = Location::any());
     RewriterVar* toBool(Location loc = Location::any());
 

src/codegen/ast_interpreter.cpp

@@ -1093,7 +1093,7 @@ Value ASTInterpreter::createFunction(AST* node, AST_arguments* args, const std::
         Value v = visit_expr(d);
         defaults.push_back(v.o);
         if (jit) {
-            defaults_var->setAttr(i++ * sizeof(void*), v);
+            defaults_var->setAttr(i++ * sizeof(void*), v, RewriterVar::SetattrType::REFUSED);
             defaults_vars.push_back(v.var);
         }
     }

src/codegen/baseline_jit.cpp

@@ -202,7 +202,7 @@ RewriterVar* JitFragmentWriter::emitCallattr(AST_expr* node, RewriterVar* obj, B
     if (args.size() > 3) {
         RewriterVar* scratch = allocate(args.size() - 3);
         for (int i = 0; i < args.size() - 3; ++i)
-            scratch->setAttr(i * sizeof(void*), args[i + 3]);
+            scratch->setAttr(i * sizeof(void*), args[i + 3], RewriterVar::SetattrType::REFUSED);
         call_args.push_back(scratch);
     } else if (keyword_names) {
         call_args.push_back(imm(0ul));
@@ -261,15 +261,17 @@ RewriterVar* JitFragmentWriter::emitCreateDict(const llvm::ArrayRef<RewriterVar*
     RewriterVar::SmallVector additional_uses;
     additional_uses.insert(additional_uses.end(), keys.begin(), keys.end());
     additional_uses.insert(additional_uses.end(), values.begin(), values.end());
-    auto rtn
-        = emitCallWithAllocatedArgs((void*)createDictHelper, { imm(keys.size()), allocArgs(keys), allocArgs(values) },
-                                    additional_uses)->setType(RefType::OWNED);
+    auto rtn = emitCallWithAllocatedArgs((void*)createDictHelper,
+                                         { imm(keys.size()), allocArgs(keys, RewriterVar::SetattrType::REFUSED),
+                                           allocArgs(values, RewriterVar::SetattrType::REFUSED) },
+                                         additional_uses)->setType(RefType::OWNED);
     for (RewriterVar* k : keys) {
-        assert(0 && "these need to be kept alive");
+        // XXX: should refConsumed also act as a use?
+        k->refUsed();
         k->refConsumed();
     }
     for (RewriterVar* v : values) {
-        assert(0 && "these need to be kept alive");
+        v->refUsed();
         v->refConsumed();
     }
     return rtn;
@@ -317,10 +319,11 @@ RewriterVar* JitFragmentWriter::emitCreateList(const llvm::ArrayRef<RewriterVar*
     if (num == 0)
         return call(false, (void*)createList)->setType(RefType::OWNED);
 
-    auto rtn = emitCallWithAllocatedArgs((void*)createListHelper, { imm(num), allocArgs(values) }, values)
-                   ->setType(RefType::OWNED);
+    auto rtn = emitCallWithAllocatedArgs((void*)createListHelper,
+                                         { imm(num), allocArgs(values, RewriterVar::SetattrType::REFUSED) },
+                                         values)->setType(RefType::OWNED);
     for (RewriterVar* v : values) {
-        assert(0 && "these need to be kept alive");
+        v->refUsed();
         v->refConsumed();
     }
     return rtn;
@@ -330,10 +333,11 @@ RewriterVar* JitFragmentWriter::emitCreateSet(const llvm::ArrayRef<RewriterVar*>
     auto num = values.size();
     if (num == 0)
         return call(false, (void*)createSet)->setType(RefType::OWNED);
-    auto rtn = emitCallWithAllocatedArgs((void*)createSetHelper, { imm(num), allocArgs(values) }, values)
-                   ->setType(RefType::OWNED);
+    auto rtn = emitCallWithAllocatedArgs((void*)createSetHelper,
+                                         { imm(num), allocArgs(values, RewriterVar::SetattrType::REFUSED) },
+                                         values)->setType(RefType::OWNED);
     for (RewriterVar* v : values) {
-        assert(0 && "these need to be kept alive");
+        v->refUsed();
         v->refConsumed();
     }
     return rtn;
@@ -355,8 +359,9 @@ RewriterVar* JitFragmentWriter::emitCreateTuple(const llvm::ArrayRef<RewriterVar
     else if (num == 3)
         r = call(false, (void*)BoxedTuple::create3, values[0], values[1], values[2])->setType(RefType::OWNED);
     else {
-        r = emitCallWithAllocatedArgs((void*)createTupleHelper, { imm(num), allocArgs(values) }, values)
-                ->setType(RefType::OWNED);
+        r = emitCallWithAllocatedArgs((void*)createTupleHelper,
+                                      { imm(num), allocArgs(values, RewriterVar::SetattrType::REFUSED) },
+                                      values)->setType(RefType::OWNED);
         for (auto a : values)
             a->refUsed();
     }
@@ -488,7 +493,7 @@ RewriterVar* JitFragmentWriter::emitRuntimeCall(AST_expr* node, RewriterVar* obj
     if (args.size() > 3) {
         RewriterVar* scratch = allocate(args.size() - 3);
         for (int i = 0; i < args.size() - 3; ++i)
-            scratch->setAttr(i * sizeof(void*), args[i + 3]);
+            scratch->setAttr(i * sizeof(void*), args[i + 3], RewriterVar::SetattrType::REFUSED);
         call_args.push_back(scratch);
     } else
         call_args.push_back(imm(0ul));
@@ -506,7 +511,7 @@ RewriterVar* JitFragmentWriter::emitRuntimeCall(AST_expr* node, RewriterVar* obj
 
     RewriterVar* args_array = nullptr;
     if (args.size()) {
-        args_array = allocArgs(args);
+        args_array = allocArgs(args, RewriterVar::SetattrType::REFUSED);
     } else
         RELEASE_ASSERT(!keyword_names_var, "0 args but keyword names are set");
 
@@ -689,7 +694,7 @@ void JitFragmentWriter::emitSetLocal(InternedString s, int vreg, bool set_closur
         v->refConsumed();
     } else {
         RewriterVar* prev = vregs_array->getAttr(8 * vreg)->setNullable(true);
-        vregs_array->setAttr(8 * vreg, v);
+        vregs_array->setAttr(8 * vreg, v, RewriterVar::SetattrType::HANDED_OFF);
         v->refConsumed();
 
         // TODO With definedness analysis, we could know whether we can skip this check (definitely defined)
@@ -838,12 +843,13 @@ bool JitFragmentWriter::finishAssembly(int continue_offset) {
 }
 
 
-RewriterVar* JitFragmentWriter::allocArgs(const llvm::ArrayRef<RewriterVar*> args) {
+RewriterVar* JitFragmentWriter::allocArgs(const llvm::ArrayRef<RewriterVar*> args,
+                                          RewriterVar::SetattrType setattr_type) {
     auto num = args.size();
     assert(num);
     RewriterVar* array = allocate(num);
     for (int i = 0; i < num; ++i)
-        array->setAttr(sizeof(void*) * i, args[i]);
+        array->setAttr(sizeof(void*) * i, args[i], setattr_type);
     return array;
 }
 

src/codegen/baseline_jit.h

@@ -281,7 +281,7 @@ public:
     bool finishAssembly(int continue_offset) override;
 
 private:
-    RewriterVar* allocArgs(const llvm::ArrayRef<RewriterVar*> args);
+    RewriterVar* allocArgs(const llvm::ArrayRef<RewriterVar*> args, RewriterVar::SetattrType);
 #ifndef NDEBUG
     std::pair<uint64_t, uint64_t> asUInt(InternedString s);
 #else

src/runtime/objmodel.cpp

@@ -1422,7 +1422,8 @@ void Box::appendNewHCAttr(BORROWED(Box*) new_attr, SetattrRewriteArgs* rewrite_a
         if (!new_array)
             r_array = rewrite_args->obj->getAttr(cls->attrs_offset + offsetof(HCAttrs, attr_list));
 
-        r_array->setAttr(numattrs * sizeof(Box*) + offsetof(HCAttrs::AttrList, attrs), rewrite_args->attrval);
+        r_array->setAttr(numattrs * sizeof(Box*) + offsetof(HCAttrs::AttrList, attrs), rewrite_args->attrval,
+                         RewriterVar::SetattrType::HANDED_OFF);
         rewrite_args->attrval->refConsumed();
 
         if (new_array)
@@ -1515,8 +1516,8 @@ void Box::setattr(BoxedString* attr, BORROWED(Box*) val, SetattrRewriteArgs* rew
                     // will tell the auto-refcount system to decref it.
                     r_hattrs->getAttr(offset * sizeof(Box*) + offsetof(HCAttrs::AttrList, attrs))
                         ->setType(RefType::OWNED);
-                    r_hattrs->setAttr(offset * sizeof(Box*) + offsetof(HCAttrs::AttrList, attrs),
-                                      rewrite_args->attrval);
+                    r_hattrs->setAttr(offset * sizeof(Box*) + offsetof(HCAttrs::AttrList, attrs), rewrite_args->attrval,
+                                      RewriterVar::SetattrType::HANDED_OFF);
                     rewrite_args->attrval->refConsumed();
 
                     rewrite_args->out_success = true;
@@ -4357,7 +4358,9 @@ void rearrangeArgumentsInternal(ParamReceiveSpec paramspec, const ParamNames* pa
                 if (varargs_idx == 2)
                     rewrite_args->arg3 = varargs_val;
                 if (varargs_idx >= 3) {
-                    rewrite_args->args->setAttr((varargs_idx - 3) * sizeof(Box*), varargs_val);
+                    rewrite_args->args->setAttr(
+                        (varargs_idx - 3) * sizeof(Box*), varargs_val,
+                        (is_owned ? RewriterVar::SetattrType::HANDED_OFF : RewriterVar::SetattrType::UNKNOWN));
                     if (is_owned) {
                         oargs_owned[varargs_idx - 3] = true;
                         varargs_val->refConsumed();
@@ -4886,13 +4889,13 @@ Box* callCLFunc(FunctionMetadata* md, CallRewriteArgs* rewrite_args, int num_out
                 RewriterVar* arg_array = rewrite_args->rewriter->allocate(4);
                 arg_vec.push_back(arg_array);
                 if (num_output_args >= 1)
-                    arg_array->setAttr(0, rewrite_args->arg1);
+                    arg_array->setAttr(0, rewrite_args->arg1, RewriterVar::SetattrType::REFUSED);
                 if (num_output_args >= 2)
-                    arg_array->setAttr(8, rewrite_args->arg2);
+                    arg_array->setAttr(8, rewrite_args->arg2, RewriterVar::SetattrType::REFUSED);
                 if (num_output_args >= 3)
-                    arg_array->setAttr(16, rewrite_args->arg3);
+                    arg_array->setAttr(16, rewrite_args->arg3, RewriterVar::SetattrType::REFUSED);
                 if (num_output_args >= 4)
-                    arg_array->setAttr(24, rewrite_args->args);
+                    arg_array->setAttr(24, rewrite_args->args, RewriterVar::SetattrType::REFUSED);
 
                 if (S == CXX)
                     rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretHelper, arg_vec)