Fix unpackIntoArray refcounting

It returns an object that keeps the returned array alive.

src/codegen/ast_interpreter.cpp

@@ -497,42 +497,52 @@ void ASTInterpreter::doStore(AST_expr* node, STOLEN(Value) value) {
         AUTO_DECREF(o.o);
         pyston::setattr(o.o, attr->attr.getBox(), value.o);
     } else if (node->type == AST_TYPE::Tuple) {
+        AUTO_DECREF(value.o);
+
         AST_Tuple* tuple = (AST_Tuple*)node;
-        Box** array = unpackIntoArray(value.o, tuple->elts.size());
+        Box* keep_alive;
+        Box** array = unpackIntoArray(value.o, tuple->elts.size(), &keep_alive);
+        AUTO_DECREF(keep_alive);
 
-        RewriterVar* array_var = NULL;
+        std::vector<RewriterVar*> array_vars;
         if (jit) {
-            array_var = jit->emitUnpackIntoArray(value, tuple->elts.size());
+            array_vars = jit->emitUnpackIntoArray(value, tuple->elts.size());
+            assert(array_vars.size() == tuple->elts.size());
         }
 
         unsigned i = 0;
         for (AST_expr* e : tuple->elts) {
-            doStore(e, Value(array[i], jit ? array_var->getAttr(i * sizeof(void*))->setType(RefType::OWNED) : NULL));
+            doStore(e, Value(array[i], jit ? array_vars[i] : NULL));
             ++i;
         }
-        Py_DECREF(value.o);
     } else if (node->type == AST_TYPE::List) {
+        AUTO_DECREF(value.o);
+
         AST_List* list = (AST_List*)node;
-        Box** array = unpackIntoArray(value.o, list->elts.size());
+        Box* keep_alive;
+        Box** array = unpackIntoArray(value.o, list->elts.size(), &keep_alive);
+        AUTO_DECREF(keep_alive);
 
-        RewriterVar* array_var = NULL;
-        if (jit)
-            array_var = jit->emitUnpackIntoArray(value, list->elts.size());
+        std::vector<RewriterVar*> array_vars;
+        if (jit) {
+            array_vars = jit->emitUnpackIntoArray(value, list->elts.size());
+            assert(array_vars.size() == list->elts.size());
+        }
 
         unsigned i = 0;
         for (AST_expr* e : list->elts) {
-            doStore(e, Value(array[i], jit ? array_var->getAttr(i * sizeof(void*)) : NULL));
+            doStore(e, Value(array[i], jit ? array_vars[i] : NULL));
             ++i;
         }
     } else if (node->type == AST_TYPE::Subscript) {
+        AUTO_DECREF(value.o);
         AST_Subscript* subscript = (AST_Subscript*)node;
 
         Value target = visit_expr(subscript->value);
-        Value slice = visit_slice(subscript->slice);
-
         AUTO_DECREF(target.o);
+
+        Value slice = visit_slice(subscript->slice);
         AUTO_DECREF(slice.o);
-        AUTO_DECREF(value.o);
 
         if (jit)
             jit->emitSetItem(target, slice, value);

src/codegen/baseline_jit.cpp

@@ -441,9 +441,21 @@ RewriterVar* JitFragmentWriter::emitUnaryop(RewriterVar* v, int op_type) {
     return emitPPCall((void*)unaryop, { v, imm(op_type) }, 2, 160)->setType(RefType::OWNED);
 }
 
-RewriterVar* JitFragmentWriter::emitUnpackIntoArray(RewriterVar* v, uint64_t num) {
-    RewriterVar* array = call(false, (void*)unpackIntoArray, v, imm(num));
-    return array;
+std::vector<RewriterVar*> JitFragmentWriter::emitUnpackIntoArray(RewriterVar* v, uint64_t num) {
+    assert(0 && "untested");
+    trap();
+
+    RewriterVar* keep_alive = allocate(1);
+
+    RewriterVar* array = call(false, (void*)unpackIntoArray, v, imm(num), keep_alive);
+
+    std::vector<RewriterVar*> rtn;
+    for (int i = 0; i < num; i++) {
+        rtn.push_back(array->getAttr(i * sizeof(void*))->setType(RefType::OWNED));
+    }
+
+    keep_alive->getAttr(0)->setType(RefType::OWNED);
+    return rtn;
 }
 
 RewriterVar* JitFragmentWriter::emitYield(RewriterVar* v) {

src/codegen/baseline_jit.h

@@ -243,7 +243,7 @@ public:
     RewriterVar* emitRuntimeCall(AST_expr* node, RewriterVar* obj, ArgPassSpec argspec,
                                  const llvm::ArrayRef<RewriterVar*> args, std::vector<BoxedString*>* keyword_names);
     RewriterVar* emitUnaryop(RewriterVar* v, int op_type);
-    RewriterVar* emitUnpackIntoArray(RewriterVar* v, uint64_t num);
+    std::vector<RewriterVar*> emitUnpackIntoArray(RewriterVar* v, uint64_t num);
     RewriterVar* emitYield(RewriterVar* v);
 
     void emitDelAttr(RewriterVar* target, BoxedString* attr);

src/codegen/compvars.cpp

@@ -496,8 +496,9 @@ public:
     }
 
     std::vector<CompilerVariable*> unpack(IREmitter& emitter, const OpInfo& info, VAR* var, int num_into) override {
-        llvm::Value* unpacked = emitter.createCall2(info.unw_info, g.funcs.unpackIntoArray, var->getValue(),
-                                                    getConstantInt(num_into, g.i64));
+        llvm::Value* scratch = emitter.getBuilder()->CreateBitCast(emitter.getScratch(sizeof(Box*)), g.llvm_value_type_ptr_ptr);
+        llvm::Value* unpacked = emitter.createCall3(info.unw_info, g.funcs.unpackIntoArray, var->getValue(),
+                                                    getConstantInt(num_into, g.i64), scratch);
         assert(unpacked->getType() == g.llvm_value_type_ptr->getPointerTo());
 
         std::vector<CompilerVariable*> rtn;
@@ -509,6 +510,8 @@ public:
 
             rtn.push_back(new ConcreteCompilerVariable(UNKNOWN, val));
         }
+
+        emitter.setType(emitter.getBuilder()->CreateLoad(scratch), RefType::OWNED);
         return rtn;
     }
 

src/core/options.cpp

@@ -27,21 +27,23 @@ int MAX_OPT_ITERATIONS = 1;
 bool LOG_IC_ASSEMBLY = 0;
 bool LOG_BJIT_ASSEMBLY = 0;
 
-bool CONTINUE_AFTER_FATAL = false;
 bool FORCE_INTERPRETER = false;
 bool FORCE_OPTIMIZE = false;
+bool ENABLE_INTERPRETER = false; // XXX
+bool ENABLE_BASELINEJIT = false; // XXX
+
+bool CONTINUE_AFTER_FATAL = false;
 bool SHOW_DISASM = false;
 bool PROFILE = false;
 bool DUMPJIT = false;
 bool TRAP = false;
 bool USE_STRIPPED_STDLIB = true; // always true
-bool ENABLE_INTERPRETER = false; // XXX
-bool ENABLE_BASELINEJIT = false; // XXX
 bool ENABLE_PYPA_PARSER = true;
 bool ENABLE_CPYTHON_PARSER = true;
 bool USE_REGALLOC_BASIC = false;
 bool PAUSE_AT_ABORT = false;
 bool ENABLE_TRACEBACKS = true;
+
 // Forces the llvm jit to use capi exceptions whenever it can, as opposed to whenever it thinks
 // it is faster.  The CALLS version is for calls that the llvm jit will make, and the THROWS version
 // is for the exceptions it will throw.

src/runtime/exceptions.cpp

@@ -206,7 +206,6 @@ extern "C" void raise3(STOLEN(Box*) arg0, STOLEN(Box*) arg1, STOLEN(Box*) arg2)
 }
 
 extern "C" void raise3_capi(STOLEN(Box*) arg0, STOLEN(Box*) arg1, STOLEN(Box*) arg2) noexcept {
-    assert(0 && "Check refcounting");
     bool reraise = arg2 != NULL && arg2 != None;
 
     ExcInfo exc_info(NULL, NULL, NULL);

src/runtime/objmodel.cpp

@@ -304,35 +304,43 @@ static void _checkUnpackingLength(i64 expected, i64 given) {
     }
 }
 
-extern "C" Box** unpackIntoArray(Box* obj, int64_t expected_size) {
-    assert(0 && "need to fix refcounting here -- need to return the Box that keeps the array alive");
+extern "C" Box** unpackIntoArray(Box* obj, int64_t expected_size, Box** out_keep_alive) {
+    if (obj->cls != tuple_cls && obj->cls != list_cls) {
+        Box* converted = PySequence_Fast(obj, "Invalid type for tuple unpacking");
+        if (!converted)
+            throwCAPIException();
+        *out_keep_alive = converted;
+        obj = converted;
+    } else {
+        *out_keep_alive = incref(obj);
+    }
+
     if (obj->cls == tuple_cls) {
         BoxedTuple* t = static_cast<BoxedTuple*>(obj);
-        _checkUnpackingLength(expected_size, t->size());
+
+        auto got_size = t->size();
+        if (expected_size != got_size)
+            Py_DECREF(*out_keep_alive);
+        _checkUnpackingLength(expected_size, got_size);
+
         for (auto e : *t)
             Py_INCREF(e);
         return &t->elts[0];
-    }
+    } else {
+        assert(obj->cls == list_cls);
 
-    if (obj->cls == list_cls) {
         BoxedList* l = static_cast<BoxedList*>(obj);
-        _checkUnpackingLength(expected_size, l->size);
+
+        auto got_size = l->size;
+        if (expected_size != got_size)
+            Py_DECREF(*out_keep_alive);
+        _checkUnpackingLength(expected_size, got_size);
+
         for (size_t i = 0; i < l->size; i++)
             Py_INCREF(l->elts->elts[i]);
         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);
-        if (elts.size() > expected_size)
-            break;
-    }
-
-    _checkUnpackingLength(expected_size, elts.size());
-    return &elts[0];
 }
 
 static void clear_slots(PyTypeObject* type, PyObject* self) noexcept {

src/runtime/objmodel.h

@@ -89,7 +89,7 @@ extern "C" Box* getclsattrMaybeNonstring(Box* obj, Box* attr);
 extern "C" Box* unaryop(Box* operand, int op_type);
 extern "C" Box* importFrom(Box* obj, BoxedString* attr);
 extern "C" Box* importStar(Box* from_module, Box* to_globals);
-extern "C" Box** unpackIntoArray(Box* obj, int64_t expected_size);
+extern "C" Box** unpackIntoArray(Box* obj, int64_t expected_size, Box** out_keep_alive);
 extern "C" void assertNameDefined(bool b, const char* name, BoxedClass* exc_cls, bool local_var_msg);
 extern "C" void assertFailDerefNameDefined(const char* name);
 extern "C" void assertFail(Box* assertion_type, Box* msg);