Fix refcounting of array-passed args

Sometimes to pass a variable number of argument we allocate an array, store
the arguments in that array, and then pass the array as the single C-level
arg.  This caused ref issues where the refcounter thought that the variables
were dead after the store -- ie before the function call -- and it would
decref them right then.  This commit adds a "refUsed" hook that allows us
to specify these sorts of "non-IR-based" uses.

Another option would have been to make the arrays fully refcounted (ie the
array itself is treated as owning references to its contents), but in this
case I think that's overkill since it's really just an ABI issue that the
refcounter didn't understand.

src/codegen/compvars.cpp

@@ -602,6 +602,8 @@ static ConcreteCompilerVariable* _call(IREmitter& emitter, const OpInfo& info, l
         llvm_args.push_back(getNullPtr(g.llvm_value_type_ptr));
     }
 
+    llvm::SmallVector<llvm::Value*, 4> array_passed_args;
+
     if (args.size() >= 4) {
         llvm::Value* arg_array;
 
@@ -616,6 +618,7 @@ static ConcreteCompilerVariable* _call(IREmitter& emitter, const OpInfo& info, l
         for (int i = 3; i < args.size(); i++) {
             llvm::Value* ptr = emitter.getBuilder()->CreateConstGEP1_32(arg_array, i - 3);
             emitter.getBuilder()->CreateStore(converted_args[i]->getValue(), ptr);
+            array_passed_args.push_back(converted_args[i]->getValue());
         }
         llvm_args.push_back(arg_array);
 
@@ -633,6 +636,7 @@ static ConcreteCompilerVariable* _call(IREmitter& emitter, const OpInfo& info, l
     //}
 
     llvm::Value* rtn;
+    llvm::Instruction* inst;
 
     // func->dump();
     // for (auto a : llvm_args)
@@ -645,7 +649,8 @@ static ConcreteCompilerVariable* _call(IREmitter& emitter, const OpInfo& info, l
 
         ICSetupInfo* pp = createCallsiteIC(info.getTypeRecorder(), args.size(), info.getBJitICInfo());
 
-        llvm::Value* uncasted = emitter.createIC(pp, func_addr, llvm_args, info.unw_info, target_exception_style);
+        llvm::Instruction* uncasted = emitter.createIC(pp, func_addr, llvm_args, info.unw_info, target_exception_style);
+        inst = uncasted;
 
         assert(llvm::cast<llvm::FunctionType>(llvm::cast<llvm::PointerType>(func->getType())->getElementType())
                    ->getReturnType() == g.llvm_value_type_ptr);
@@ -659,7 +664,8 @@ static ConcreteCompilerVariable* _call(IREmitter& emitter, const OpInfo& info, l
         //}
         // printf("%ld %ld\n", llvm_args.size(), args.size());
         // printf("\n");
-        rtn = emitter.createCall(info.unw_info, func, llvm_args, target_exception_style);
+        inst = emitter.createCall(info.unw_info, func, llvm_args, target_exception_style);
+        rtn = inst;
     }
 
     if (rtn_type->getBoxType() == rtn_type)
@@ -670,6 +676,9 @@ static ConcreteCompilerVariable* _call(IREmitter& emitter, const OpInfo& info, l
         emitter.checkAndPropagateCapiException(info.unw_info, rtn, getNullPtr(g.llvm_value_type_ptr));
     }
 
+    for (auto v : array_passed_args)
+        emitter.refUsed(v, inst);
+
     return new ConcreteCompilerVariable(rtn_type, rtn);
 }
 
@@ -815,6 +824,8 @@ CompilerVariable* makeFunction(IREmitter& emitter, FunctionMetadata* f, Compiler
         emitter.setType(closure_v, RefType::BORROWED);
     }
 
+    llvm::SmallVector<llvm::Value*, 4> array_passed_args;
+
     llvm::Value* scratch;
     if (defaults.size()) {
         scratch = emitter.getScratch(defaults.size() * sizeof(Box*));
@@ -822,6 +833,7 @@ CompilerVariable* makeFunction(IREmitter& emitter, FunctionMetadata* f, Compiler
         int i = 0;
         for (auto d : defaults) {
             llvm::Value* v = d->getValue();
+            array_passed_args.push_back(v);
             llvm::Value* p = emitter.getBuilder()->CreateConstGEP1_32(scratch, i);
             emitter.getBuilder()->CreateStore(v, p);
             i++;
@@ -834,12 +846,17 @@ CompilerVariable* makeFunction(IREmitter& emitter, FunctionMetadata* f, Compiler
 
     // We know this function call can't throw, so it's safe to use emitter.getBuilder()->CreateCall() rather than
     // emitter.createCall().
-    llvm::Value* boxed = emitter.getBuilder()->CreateCall(
+    llvm::Instruction* boxed = emitter.getBuilder()->CreateCall(
         g.funcs.createFunctionFromMetadata,
         std::vector<llvm::Value*>{ embedRelocatablePtr(f, g.llvm_functionmetadata_type_ptr), closure_v, globals,
                                    scratch, getConstantInt(defaults.size(), g.i64) });
     emitter.setType(boxed, RefType::OWNED);
 
+    // The refcounter needs to know that this call "uses" the arguments that got passed via scratch.
+    for (auto v : array_passed_args) {
+        emitter.refUsed(v, boxed);
+    }
+
     return new ConcreteCompilerVariable(typeFromClass(function_cls), boxed);
 }
 
@@ -2467,6 +2484,8 @@ public:
         llvm::Value* _scratch = emitter.getScratch(v.size() * sizeof(void*));
         auto scratch = emitter.getBuilder()->CreateBitCast(_scratch, g.llvm_value_type_ptr->getPointerTo());
 
+        llvm::SmallVector<llvm::Value*, 4> array_passed_args;
+
         // First, convert all the args, before putting any in the scratch.
         // Do it this way in case any of the conversions themselves need scratch space
         // (ie nested tuples).
@@ -2481,11 +2500,15 @@ public:
         for (int i = 0; i < v.size(); i++) {
             llvm::Value* ptr = emitter.getBuilder()->CreateConstGEP1_32(scratch, i);
             emitter.getBuilder()->CreateStore(converted_args[i]->getValue(), ptr);
+            array_passed_args.push_back(converted_args[i]->getValue());
         }
 
-        llvm::Value* rtn = emitter.getBuilder()->CreateCall2(g.funcs.createTuple, nelts, scratch);
+        llvm::Instruction* rtn = emitter.getBuilder()->CreateCall2(g.funcs.createTuple, nelts, scratch);
         emitter.setType(rtn, RefType::OWNED);
 
+        for (auto v : array_passed_args)
+            emitter.refUsed(v, rtn);
+
         return new ConcreteCompilerVariable(other_type, rtn);
     }
 

src/codegen/irgen.h

@@ -111,6 +111,7 @@ public:
 
     virtual llvm::Value* setType(llvm::Value* v, RefType reftype) = 0;
     virtual void refConsumed(llvm::Value* v, llvm::Instruction* inst) = 0;
+    virtual void refUsed(llvm::Value* v, llvm::Instruction* inst) = 0;
     virtual ConcreteCompilerVariable* getNone() = 0;
 };
 
@@ -203,12 +204,14 @@ private:
         //llvm::SmallVector<llvm::Instruction*, 2> ref_consumers;
     };
     llvm::DenseMap<llvm::Instruction*, llvm::SmallVector<llvm::Value*, 4>> refs_consumed;
+    llvm::DenseMap<llvm::Instruction*, llvm::SmallVector<llvm::Value*, 4>> refs_used;
     llvm::ValueMap<llvm::Value*, RefcountState> vars;
 public:
 
     llvm::Value* setType(llvm::Value* v, RefType reftype);
     llvm::Value* setNullable(llvm::Value* v, bool nullable = true);
     void refConsumed(llvm::Value* v, llvm::Instruction*);
+    void refUsed(llvm::Value* v, llvm::Instruction*);
     static void addRefcounts(IRGenState* state);
 };
 

src/codegen/irgen/irgenerator.cpp

@@ -650,6 +650,10 @@ public:
         irstate->getRefcounts()->refConsumed(v, inst);
     }
 
+    void refUsed(llvm::Value* v, llvm::Instruction* inst) override {
+        irstate->getRefcounts()->refUsed(v, inst);
+    }
+
     llvm::Value* setType(llvm::Value* v, RefType reftype) override {
         assert(llvm::isa<PointerType>(v->getType()));
 

src/codegen/irgen/refcounts.cpp

@@ -73,7 +73,12 @@ void RefcountTracker::refConsumed(llvm::Value* v, llvm::Instruction* inst) {
     assert(this->vars[v].reftype != RefType::UNKNOWN);
 
     this->refs_consumed[inst].push_back(v);
-    //var.ref_consumers.push_back(inst);
+}
+
+void RefcountTracker::refUsed(llvm::Value* v, llvm::Instruction* inst) {
+    assert(this->vars[v].reftype != RefType::UNKNOWN);
+
+    this->refs_used[inst].push_back(v);
 }
 
 void remapPhis(llvm::BasicBlock* in_block, llvm::BasicBlock* from_block, llvm::BasicBlock* new_from_block) {
@@ -672,6 +677,11 @@ void RefcountTracker::addRefcounts(IRGenState* irstate) {
                 num_times_as_op[v]; // just make sure it appears in there
             }
 
+            for (auto v : rt->refs_used[&I]) {
+                assert(rt->vars[v].reftype != RefType::UNKNOWN);
+                num_times_as_op[v]++;
+            }
+
             for (llvm::Value* op : I.operands()) {
                 auto it = rt->vars.find(op);
                 if (it == rt->vars.end())

test/tests/scratch_refcounting.py

+# Test various ways that we use "scratch" space.
+# Tests some regressions where we weren't correctly counting those as reference uses.
+
+def f(arg=2.0**4):
+    print arg
+f()
+
+def f2(a, b, c, d):
+    print a
+    print b
+    print c
+    print d
+f2(2.0**5, 2.0**6, 2.0**7, 2.0**8)
+
+def f3():
+    t = (2.0**9, 2.0**10, 2.0**11, 2.0**12, 2.0**13)
+    print t
+f3()