Fix a bunch of exception handling

Makefile

@@ -397,7 +397,7 @@ all: pyston_dbg pyston_release pyston_gcc unittests check-deps $(RUN_DEPS)
 ALL_HEADERS := $(wildcard src/*/*.h) $(wildcard src/*/*/*.h) $(wildcard from_cpython/Include/*.h)
 tags: $(SRCS) $(OPTIONAL_SRCS) $(FROM_CPYTHON_SRCS) $(ALL_HEADERS)
 	$(ECHO) Calculating tags...
-	$(VERB) $(CTAGS) -I noexcept -I noexcept+ -I PYSTON_NOEXCEPT $^
+	$(VERB) $(CTAGS) -I noexcept -I noexcept+ -I PYSTON_NOEXCEPT -I STOLEN -I BORROWED $^
 
 TAGS: $(SRCS) $(OPTIONAL_SRCS) $(FROM_CPYTHON_SRCS) $(ALL_HEADERS)
 	$(ECHO) Calculating TAGS...

docs/REFCOUNTING.md

@@ -19,6 +19,9 @@ AUTO_DECREF, autoDecref, incref
 
 Invariant: all paths through a function should leave each variable with 0 net refs.  (This includes paths out via exceptions.)
 
+Special cases:
+- Most exception-related functions steal refs.  If you throw an exception in C code (via `throw [ExcInfo] e`), it consumes the refs.
+
 ## Refcounting in the rewriter, baseline jit, and llvm jit
 
 ## Debugging

from_cpython/Python/errors.c

@@ -189,6 +189,7 @@ PyErr_NormalizeException(PyObject** exc, PyObject** val, PyObject** tb)
 
             if (res == NULL)
                 goto finally;
+            Py_DECREF(value);
             value = res;
         }
         /* if the class of the instance doesn't exactly match the

src/codegen/ast_interpreter.cpp

@@ -936,6 +936,8 @@ Value ASTInterpreter::visit_langPrimitive(AST_LangPrimitive* node) {
     } else if (node->opcode == AST_LangPrimitive::SET_EXC_INFO) {
         assert(node->args.size() == 3);
 
+        assert(0 && "check refcounting -- call setFrameExcInfo instead of setting it directly");
+
         Value type = visit_expr(node->args[0]);
         assert(type.o);
         Value value = visit_expr(node->args[1]);

src/codegen/irgen/irgenerator.cpp

@@ -882,9 +882,7 @@ private:
                     // trigger an exception, but the irgenerator will know that definitely-defined
                     // local symbols will not throw.
                     emitter.getBuilder()->CreateUnreachable();
-                    exc_type = undefVariable();
-                    exc_value = undefVariable();
-                    exc_tb = undefVariable();
+                    exc_type = exc_value = exc_tb = emitter.getNone();
                     endBlock(DEAD);
                 }
 
@@ -993,16 +991,17 @@ private:
                 auto* builder = emitter.getBuilder();
 
                 llvm::Value* frame_info = irstate->getFrameInfoVar();
-                llvm::Value* exc_info = builder->CreateConstInBoundsGEP2_32(frame_info, 0, 0);
-                assert(exc_info->getType() == g.llvm_excinfo_type->getPointerTo());
 
                 ConcreteCompilerVariable* converted_type = type->makeConverted(emitter, UNKNOWN);
-                builder->CreateStore(converted_type->getValue(), builder->CreateConstInBoundsGEP2_32(exc_info, 0, 0));
                 ConcreteCompilerVariable* converted_value = value->makeConverted(emitter, UNKNOWN);
-                builder->CreateStore(converted_value->getValue(), builder->CreateConstInBoundsGEP2_32(exc_info, 0, 1));
                 ConcreteCompilerVariable* converted_traceback = traceback->makeConverted(emitter, UNKNOWN);
-                builder->CreateStore(converted_traceback->getValue(),
-                                     builder->CreateConstInBoundsGEP2_32(exc_info, 0, 2));
+
+                auto inst = emitter.createCall(UnwindInfo::cantUnwind(), g.funcs.setFrameExcInfo,
+                                               { frame_info, converted_type->getValue(), converted_value->getValue(),
+                                                 converted_traceback->getValue() });
+                emitter.refConsumed(converted_type->getValue(), inst);
+                emitter.refConsumed(converted_value->getValue(), inst);
+                emitter.refConsumed(converted_traceback->getValue(), inst);
 
                 return emitter.getNone();
             }
@@ -2421,13 +2420,19 @@ private:
             }
         }
 
+        llvm::Instruction* inst;
         if (target_exception_style == CAPI) {
-            emitter.createCall(unw_info, g.funcs.raise3_capi, args, CAPI);
+            inst = emitter.createCall(unw_info, g.funcs.raise3_capi, args, CAPI);
             emitter.checkAndPropagateCapiException(unw_info, getNullPtr(g.llvm_value_type_ptr),
                                                    getNullPtr(g.llvm_value_type_ptr));
         } else {
-            emitter.createCall(unw_info, g.funcs.raise3, args, CXX);
+            inst = emitter.createCall(unw_info, g.funcs.raise3, args, CXX);
         }
+
+        for (auto a : args) {
+            emitter.refConsumed(a, inst);
+        }
+
         emitter.getBuilder()->CreateUnreachable();
 
         endBlock(DEAD);
@@ -2963,7 +2968,6 @@ public:
                 llvm::Value* exc_value = emitter.getBuilder()->CreateLoad(exc_value_ptr);
                 llvm::Value* exc_traceback = emitter.getBuilder()->CreateLoad(exc_traceback_ptr);
 
-                // TODO: nullable?
                 emitter.setType(exc_type, RefType::OWNED);
                 emitter.setType(exc_value, RefType::OWNED);
                 emitter.setType(exc_traceback, RefType::OWNED);

src/codegen/irgen/refcounts.cpp

@@ -90,12 +90,15 @@ void remapPhis(llvm::BasicBlock* in_block, llvm::BasicBlock* from_block, llvm::B
     }
 }
 
+typedef llvm::DenseMap<std::pair<llvm::BasicBlock*, llvm::BasicBlock*>, llvm::Instruction*> InsertionCache;
 llvm::Instruction* findInsertionPoint(llvm::BasicBlock* BB, llvm::BasicBlock* from_bb,
-                                      llvm::DenseMap<llvm::BasicBlock*, llvm::Instruction*> cache) {
+                                      InsertionCache& cache) {
     assert(BB);
     assert(BB != from_bb);
 
-    auto it = cache.find(BB);
+    auto key = std::make_pair(BB, from_bb);
+
+    auto it = cache.find(key);
     if (it != cache.end())
         return it->second;
 
@@ -105,6 +108,8 @@ llvm::Instruction* findInsertionPoint(llvm::BasicBlock* BB, llvm::BasicBlock* fr
 
         llvm::BasicBlock* breaker_block = llvm::BasicBlock::Create(g.context, "breaker", from_bb->getParent(), BB);
         llvm::BranchInst::Create(BB, breaker_block);
+        // llvm::outs() << "Breaking edge from " << from_bb->getName() << " to " << BB->getName() << "; name is "
+        //<< breaker_block->getName() << '\n';
 
         auto terminator = from_bb->getTerminator();
 
@@ -124,8 +129,8 @@ llvm::Instruction* findInsertionPoint(llvm::BasicBlock* BB, llvm::BasicBlock* fr
 
         remapPhis(BB, from_bb, breaker_block);
 
-        cache[BB] = breaker_block->getFirstInsertionPt();
-        return cache[BB];
+        cache[key] = breaker_block->getFirstInsertionPt();
+        return cache[key];
     }
 
     if (llvm::isa<llvm::LandingPadInst>(*BB->begin())) {
@@ -134,12 +139,12 @@ llvm::Instruction* findInsertionPoint(llvm::BasicBlock* BB, llvm::BasicBlock* fr
         ++it;
         ++it;
         ++it;
-        cache[BB] = it;
+        cache[key] = it;
         return &*it;
     } else {
         for (llvm::Instruction& I : *BB) {
             if (!llvm::isa<llvm::PHINode>(I) && !llvm::isa<llvm::AllocaInst>(I)) {
-                cache[BB] = &I;
+                cache[key] = &I;
                 return &I;
             }
         }
@@ -782,9 +787,32 @@ void RefcountTracker::addRefcounts(IRGenState* irstate) {
 
     ASSERT(states.size() == f->size(), "We didn't process all nodes...");
 
-    llvm::DenseMap<llvm::BasicBlock*, llvm::Instruction*> insertion_pts;
-    for (auto&& p : states) {
-        auto&& state = p.second;
+    // Note: we iterate over the basicblocks in the order they appear in the llvm function;
+    // iterating over states directly is nondeterministic.
+    std::vector<llvm::BasicBlock*> basic_blocks;
+    for (auto&& BB : *f)
+        basic_blocks.push_back(&BB);
+
+    // First, find all insertion points.  This may change the CFG by breaking critical edges.
+    InsertionCache insertion_pts;
+    for (auto bb : basic_blocks) {
+        auto&& state = states[bb];
+        for (auto& op : state.increfs) {
+            auto insertion_pt = op.insertion_inst;
+            if (!insertion_pt)
+                insertion_pt = findInsertionPoint(op.insertion_bb, op.insertion_from_bb, insertion_pts);
+        }
+        for (auto& op : state.decrefs) {
+            auto insertion_pt = op.insertion_inst;
+            if (!insertion_pt)
+                insertion_pt = findInsertionPoint(op.insertion_bb, op.insertion_from_bb, insertion_pts);
+        }
+    }
+
+    // Then use the insertion points (it's the same code but this time it will hit the cache).
+    // This may change the CFG by adding decref's
+    for (auto bb : basic_blocks) {
+        auto&& state = states[bb];
         for (auto& op : state.increfs) {
             auto insertion_pt = op.insertion_inst;
             if (!insertion_pt)

src/codegen/runtime_hooks.cpp

@@ -199,6 +199,7 @@ void initGlobalFuncs(GlobalState& g) {
     GET(initFrame);
     GET(deinitFrame);
     GET(makePendingCalls);
+    GET(setFrameExcInfo);
 
     GET(getattr);
     GET(getattr_capi);

src/codegen/runtime_hooks.h

@@ -34,7 +34,8 @@ struct GlobalFuncs {
 
     llvm::Value* boxInt, *unboxInt, *boxFloat, *unboxFloat, *createFunctionFromMetadata, *getFunctionMetadata,
         *boxInstanceMethod, *boxBool, *unboxBool, *createTuple, *createDict, *createList, *createSlice,
-        *createUserClass, *createClosure, *createGenerator, *createSet, *initFrame, *deinitFrame, *makePendingCalls;
+        *createUserClass, *createClosure, *createGenerator, *createSet, *initFrame, *deinitFrame, *makePendingCalls,
+        *setFrameExcInfo;
     llvm::Value* getattr, *getattr_capi, *setattr, *delattr, *delitem, *delGlobal, *nonzero, *binop, *compare,
         *augbinop, *unboxedLen, *getitem, *getitem_capi, *getclsattr, *getGlobal, *setitem, *unaryop, *import,
         *importFrom, *importStar, *repr, *exceptionMatches, *yield, *getiterHelper, *hasnext, *setGlobal, *apply_slice;

src/runtime/capi.cpp

@@ -1786,6 +1786,7 @@ Box* BoxedCApiFunction::tppCall(Box* _self, CallRewriteArgs* rewrite_args, ArgPa
                     = rewrite_args->rewriter->call(true, (void*)BaseException->tp_new, rewrite_args->arg1,
                                                    rewrite_args->rewriter->loadConst(0, Location::forArg(1)),
                                                    rewrite_args->rewriter->loadConst(0, Location::forArg(2)));
+                rewrite_args->out_rtn->setType(RefType::OWNED);
                 rewrite_args->out_success = true;
             }
             return rtn;

src/runtime/exceptions.cpp

@@ -92,13 +92,17 @@ bool ExcInfo::matches(BoxedClass* cls) const {
 }
 
 // takes the three arguments of a `raise' and produces the ExcInfo to throw
-ExcInfo excInfoForRaise(Box* type, Box* value, Box* tb) {
+ExcInfo excInfoForRaise(STOLEN(Box*) type, STOLEN(Box*) value, STOLEN(Box*) tb) {
     assert(type && value && tb); // use None for default behavior, not nullptr
     // TODO switch this to PyErr_Normalize
 
     if (tb == None) {
+        Py_DECREF(tb);
         tb = NULL;
     } else if (tb != NULL && !PyTraceBack_Check(tb)) {
+        Py_DECREF(type);
+        Py_DECREF(value);
+        Py_XDECREF(tb);
         raiseExcHelper(TypeError, "raise: arg 3 must be a traceback or None");
     }
 
@@ -115,6 +119,9 @@ ExcInfo excInfoForRaise(Box* type, Box* value, Box* tb) {
         PyErr_NormalizeException(&type, &value, &tb);
 
         if (!PyExceptionInstance_Check(value)) {
+            Py_DECREF(type);
+            Py_DECREF(value);
+            Py_XDECREF(tb);
             raiseExcHelper(TypeError, "calling %s() should have returned an instance of "
                                       "BaseException, not '%s'",
                            ((PyTypeObject*)type)->tp_name, Py_TYPE(value)->tp_name);
@@ -122,6 +129,9 @@ ExcInfo excInfoForRaise(Box* type, Box* value, Box* tb) {
     } else if (PyExceptionInstance_Check(type)) {
         /* Raising an instance.  The value should be a dummy. */
         if (value != Py_None) {
+            Py_DECREF(type);
+            Py_DECREF(value);
+            Py_XDECREF(tb);
             raiseExcHelper(TypeError, "instance exception may not have a separate value");
         } else {
             /* Normalize to raise <class>, <instance> */
@@ -131,6 +141,9 @@ ExcInfo excInfoForRaise(Box* type, Box* value, Box* tb) {
             Py_INCREF(type);
         }
     } else {
+        Py_DECREF(type);
+        Py_DECREF(value);
+        Py_XDECREF(tb);
         /* Not something you can raise.  You get an exception
            anyway, just not what you specified :-) */
         raiseExcHelper(TypeError, "exceptions must be old-style classes or "
@@ -141,7 +154,7 @@ ExcInfo excInfoForRaise(Box* type, Box* value, Box* tb) {
     assert(PyExceptionClass_Check(type));
 
     if (tb == NULL) {
-        tb = None;
+        tb = incref(None);
     }
 
     return ExcInfo(type, value, tb);
@@ -178,7 +191,7 @@ extern "C" void raise0_capi(ExcInfo* frame_exc_info) noexcept {
     PyErr_Restore(frame_exc_info->type, frame_exc_info->value, frame_exc_info->traceback);
 }
 
-extern "C" void raise3(Box* arg0, Box* arg1, Box* arg2) {
+extern "C" void raise3(STOLEN(Box*) arg0, STOLEN(Box*) arg1, STOLEN(Box*) arg2) {
     bool reraise = arg2 != NULL && arg2 != None;
     auto exc_info = excInfoForRaise(arg0, arg1, arg2);
 
@@ -189,7 +202,7 @@ extern "C" void raise3(Box* arg0, Box* arg1, Box* arg2) {
     throw exc_info;
 }
 
-extern "C" void raise3_capi(Box* arg0, Box* arg1, Box* arg2) noexcept {
+extern "C" void raise3_capi(STOLEN(Box*) arg0, STOLEN(Box*) arg1, STOLEN(Box*) arg2) noexcept {
     bool reraise = arg2 != NULL && arg2 != None;
 
     ExcInfo exc_info(NULL, NULL, NULL);
@@ -262,9 +275,9 @@ extern "C" void caughtCapiException() {
 extern "C" void reraiseCapiExcAsCxx() {
     ensureCAPIExceptionSet();
     // TODO: we are normalizing to many times?
-    ExcInfo e = excInfoForRaise(cur_thread_state.curexc_type, cur_thread_state.curexc_value,
-                                cur_thread_state.curexc_traceback);
-    PyErr_Clear();
+    Box* type, *value, *tb;
+    PyErr_Fetch(&type, &value, &tb);
+    ExcInfo e = excInfoForRaise(type, value, tb);
     startReraise();
     throw e;
 }

src/runtime/frame.cpp

@@ -217,6 +217,23 @@ extern "C" void deinitFrame(FrameInfo* frame_info) {
         Py_CLEAR(frame_info->frame_obj);
     }
     Py_CLEAR(frame_info->boxedLocals);
+
+    if (frame_info->exc.type) {
+        Py_CLEAR(frame_info->exc.type);
+        Py_CLEAR(frame_info->exc.value);
+        Py_CLEAR(frame_info->exc.traceback);
+    }
+}
+
+extern "C" void setFrameExcInfo(FrameInfo* frame_info, STOLEN(Box*) type, STOLEN(Box*) value, STOLEN(Box*) tb) {
+    if (frame_info->exc.type) {
+        Py_DECREF(frame_info->exc.type);
+        Py_DECREF(frame_info->exc.value);
+        Py_DECREF(frame_info->exc.traceback);
+    }
+    frame_info->exc.type = type;
+    frame_info->exc.value = value;
+    frame_info->exc.traceback = tb;
 }
 
 extern "C" int PyFrame_GetLineNumber(PyFrameObject* _f) noexcept {

src/runtime/inline/link_forcer.cpp

@@ -74,6 +74,7 @@ void force() {
     FORCE(initFrame);
     FORCE(deinitFrame);
     FORCE(makePendingCalls);
+    FORCE(setFrameExcInfo);
 
     FORCE(getattr);
     FORCE(getattr_capi);

src/runtime/types.h

@@ -1270,6 +1270,7 @@ Box* getFrame(int depth);
 void frameInvalidateBack(BoxedFrame* frame);
 extern "C" void deinitFrame(FrameInfo* frame_info);
 extern "C" void initFrame(FrameInfo* frame_info);
+extern "C" void setFrameExcInfo(FrameInfo* frame_info, STOLEN(Box*) type, STOLEN(Box*) value, STOLEN(Box*) tb);
 
 inline BoxedString* boxString(llvm::StringRef s) {
     if (s.size() <= 1) {