Fix some bjit/rewriter crashes

Need a 'refUsed()' function like we have in the llvm jit, for when we pass args
through an args array.

src/asm_writing/rewriter.cpp

@@ -1294,6 +1294,11 @@ void RewriterVar::refConsumed(RewriterAction* action) {
     action->consumed_refs.emplace_back(this);
 }
 
+void RewriterVar::refUsed() {
+    // TODO: This is a pretty silly implementation that might prevent other optimizations?
+    rewriter->addAction([=]() { this->bumpUse(); }, { this }, ActionType::NORMAL);
+}
+
 bool RewriterVar::needsDecref() {
     return reftype == RefType::OWNED && !this->refHandedOff();
 }

src/asm_writing/rewriter.h

@@ -168,6 +168,8 @@ public:
     // if no action is specified it will assume the last action consumed the reference
     void refConsumed(RewriterAction* action = NULL);
 
+    void refUsed();
+
 
     template <typename Src, typename Dst> inline RewriterVar* getAttrCast(int offset, Location loc = Location::any());
 

src/codegen/ast_interpreter.cpp

@@ -1083,6 +1083,7 @@ Value ASTInterpreter::createFunction(AST* node, AST_arguments* args, const std::
     FunctionMetadata* md = wrapFunction(node, args, body, source_info);
 
     std::vector<Box*> defaults;
+    llvm::SmallVector<RewriterVar*, 4> defaults_vars;
 
     RewriterVar* defaults_var = NULL;
     if (jit)
@@ -1091,8 +1092,10 @@ Value ASTInterpreter::createFunction(AST* node, AST_arguments* args, const std::
     for (AST_expr* d : args->defaults) {
         Value v = visit_expr(d);
         defaults.push_back(v.o);
-        if (jit)
+        if (jit) {
             defaults_var->setAttr(i++ * sizeof(void*), v);
+            defaults_vars.push_back(v.var);
+        }
     }
     defaults.push_back(0);
     AUTO_XDECREF_ARRAY(defaults.data(), defaults.size());
@@ -1156,6 +1159,10 @@ Value ASTInterpreter::createFunction(AST* node, AST_arguments* args, const std::
             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()))->setType(RefType::OWNED);
+
+        for (auto d_var : defaults_vars) {
+            d_var->refUsed();
+        }
     }
 
     rtn.o = createFunctionFromMetadata(md, closure, passed_globals, u.il);

src/codegen/baseline_jit.cpp

@@ -211,7 +211,11 @@ RewriterVar* JitFragmentWriter::emitCallattr(AST_expr* node, RewriterVar* obj, B
     if (keyword_names)
         call_args.push_back(imm(keyword_names));
 
-    return emitPPCall((void*)callattr, call_args, 2, 640, node, type_recorder).first->setType(RefType::OWNED);
+    auto r = emitPPCall((void*)callattr, call_args, 2, 640, node, type_recorder).first->setType(RefType::OWNED);
+    for (int i = 3; i < args.size(); i++) {
+        args[i]->refUsed();
+    }
+    return r;
 #else
     // We could make this faster but for now: keep it simple, stupid...
     RewriterVar* attr_var = imm(attr);
@@ -236,7 +240,11 @@ RewriterVar* JitFragmentWriter::emitCallattr(AST_expr* node, RewriterVar* obj, B
     if (keyword_names_var)
         call_args.push_back(keyword_names_var);
 
-    return emitCallWithAllocatedArgs((void*)callattrHelper, call_args, args)->setType(RefType::OWNED);
+    auto r = emitCallWithAllocatedArgs((void*)callattrHelper, call_args, args)->setType(RefType::OWNED);
+    for (int i = 0; i < args.size(); i++) {
+        args[i]->refUsed();
+    }
+    return r;
 #endif
 }
 
@@ -257,9 +265,11 @@ RewriterVar* JitFragmentWriter::emitCreateDict(const llvm::ArrayRef<RewriterVar*
         = emitCallWithAllocatedArgs((void*)createDictHelper, { imm(keys.size()), allocArgs(keys), allocArgs(values) },
                                     additional_uses)->setType(RefType::OWNED);
     for (RewriterVar* k : keys) {
+        assert(0 && "these need to be kept alive");
         k->refConsumed();
     }
     for (RewriterVar* v : values) {
+        assert(0 && "these need to be kept alive");
         v->refConsumed();
     }
     return rtn;
@@ -310,6 +320,7 @@ RewriterVar* JitFragmentWriter::emitCreateList(const llvm::ArrayRef<RewriterVar*
     auto rtn = emitCallWithAllocatedArgs((void*)createListHelper, { imm(num), allocArgs(values) }, values)
                    ->setType(RefType::OWNED);
     for (RewriterVar* v : values) {
+        assert(0 && "these need to be kept alive");
         v->refConsumed();
     }
     return rtn;
@@ -322,6 +333,7 @@ RewriterVar* JitFragmentWriter::emitCreateSet(const llvm::ArrayRef<RewriterVar*>
     auto rtn = emitCallWithAllocatedArgs((void*)createSetHelper, { imm(num), allocArgs(values) }, values)
                    ->setType(RefType::OWNED);
     for (RewriterVar* v : values) {
+        assert(0 && "these need to be kept alive");
         v->refConsumed();
     }
     return rtn;
@@ -342,9 +354,12 @@ RewriterVar* JitFragmentWriter::emitCreateTuple(const llvm::ArrayRef<RewriterVar
         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])->setType(RefType::OWNED);
-    else
+    else {
         r = emitCallWithAllocatedArgs((void*)createTupleHelper, { imm(num), allocArgs(values) }, values)
                 ->setType(RefType::OWNED);
+        for (auto a : values)
+            a->refUsed();
+    }
     return r;
 }
 
@@ -480,7 +495,11 @@ RewriterVar* JitFragmentWriter::emitRuntimeCall(AST_expr* node, RewriterVar* obj
     if (keyword_names)
         call_args.push_back(imm(keyword_names));
 
-    return emitPPCall((void*)runtimeCall, call_args, 2, 640, node, type_recorder).first->setType(RefType::OWNED);
+    auto r = emitPPCall((void*)runtimeCall, call_args, 2, 640, node, type_recorder).first->setType(RefType::OWNED);
+    for (int i = 3; i < args.size(); i++) {
+        args[i]->refUsed();
+    }
+    return r;
 #else
     RewriterVar* argspec_var = imm(argspec.asInt());
     RewriterVar* keyword_names_var = keyword_names ? imm(keyword_names) : nullptr;
@@ -500,7 +519,11 @@ RewriterVar* JitFragmentWriter::emitRuntimeCall(AST_expr* node, RewriterVar* obj
     if (keyword_names_var)
         call_args.push_back(keyword_names_var);
 
-    return emitCallWithAllocatedArgs((void*)runtimeCallHelper, call_args, args)->setType(RefType::OWNED);
+    auto r = emitCallWithAllocatedArgs((void*)runtimeCallHelper, call_args, args)->setType(RefType::OWNED);
+    for (int i = 0; i < args.size(); i++) {
+        args[i]->refUsed();
+    }
+    return r;
 #endif
 }
 

src/runtime/objmodel.cpp

@@ -4900,6 +4900,15 @@ Box* callCLFunc(FunctionMetadata* md, CallRewriteArgs* rewrite_args, int num_out
                 else
                     rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretHelperCapi, arg_vec)
                                                 ->setType(RefType::OWNED);
+
+                if (num_output_args >= 1)
+                    rewrite_args->arg1->refUsed();
+                if (num_output_args >= 2)
+                    rewrite_args->arg2->refUsed();
+                if (num_output_args >= 3)
+                    rewrite_args->arg3->refUsed();
+                if (num_output_args >= 4)
+                    rewrite_args->args->refUsed();
             }
 
             rewrite_args->out_success = true;