Support checking for return codes other than NULL

src/asm_writing/assembler.cpp

@@ -712,22 +712,7 @@ void Assembler::cmp(Register reg1, Register reg2) {
 }
 
 void Assembler::cmp(Register reg, Immediate imm) {
-    int64_t val = imm.val;
-    assert((-1L << 31) <= val && val < (1L << 31) - 1);
-
-    int reg_idx = reg.regnum;
-
-    int rex = REX_W;
-    if (reg_idx >= 8) {
-        rex |= REX_B;
-        reg_idx -= 8;
-    }
-    assert(0 <= reg_idx && reg_idx < 8);
-
-    emitRex(rex);
-    emitByte(0x81);
-    emitModRM(0b11, 7, reg_idx);
-    emitInt(val, 4);
+    emitArith(imm, reg, OPCODE_CMP);
 }
 
 void Assembler::cmp(Indirect mem, Immediate imm) {

src/asm_writing/assembler.h

@@ -73,7 +73,7 @@ private:
     uint8_t* addr;
     bool failed; // if the rewrite failed at the assembly-generation level for some reason
 
-    static const uint8_t OPCODE_ADD = 0b000, OPCODE_SUB = 0b101;
+    static const uint8_t OPCODE_ADD = 0b000, OPCODE_SUB = 0b101, OPCODE_CMP = 0b111;
     static const uint8_t REX_B = 1, REX_X = 2, REX_R = 4, REX_W = 8;
 
 #ifndef NDEBUG

src/asm_writing/rewriter.cpp

@@ -1523,17 +1523,20 @@ void Rewriter::_allocateAndCopyPlus1(RewriterVar* result, RewriterVar* first_ele
     assertConsistent();
 }
 
-void Rewriter::checkAndThrowCAPIException(RewriterVar* r) {
+void Rewriter::checkAndThrowCAPIException(RewriterVar* r, int64_t exc_val) {
     STAT_TIMER(t0, "us_timer_rewriter", 10);
 
-    addAction([=]() { this->_checkAndThrowCAPIException(r); }, { r }, ActionType::MUTATION);
+    addAction([=]() { this->_checkAndThrowCAPIException(r, exc_val); }, { r }, ActionType::MUTATION);
 }
 
-void Rewriter::_checkAndThrowCAPIException(RewriterVar* r) {
+void Rewriter::_checkAndThrowCAPIException(RewriterVar* r, int64_t exc_val) {
     assembler->comment("_checkAndThrowCAPIException");
 
     assembler::Register var_reg = r->getInReg();
-    assembler->test(var_reg, var_reg);
+    if (exc_val == 0)
+        assembler->test(var_reg, var_reg);
+    else
+        assembler->cmp(var_reg, assembler::Immediate(exc_val));
 
     {
         assembler::ForwardJump jnz(*assembler, assembler::COND_NOT_ZERO);

src/asm_writing/rewriter.h

@@ -430,7 +430,7 @@ protected:
     int _allocate(RewriterVar* result, int n);
     void _allocateAndCopy(RewriterVar* result, RewriterVar* array, int n);
     void _allocateAndCopyPlus1(RewriterVar* result, RewriterVar* first_elem, RewriterVar* rest, int n_rest);
-    void _checkAndThrowCAPIException(RewriterVar* r);
+    void _checkAndThrowCAPIException(RewriterVar* r, int64_t exc_val);
 
     // The public versions of these are in RewriterVar
     void _addGuard(RewriterVar* var, RewriterVar* val_constant);
@@ -515,8 +515,8 @@ public:
     RewriterVar* allocateAndCopy(RewriterVar* array, int n);
     RewriterVar* allocateAndCopyPlus1(RewriterVar* first_elem, RewriterVar* rest, int n_rest);
 
-    // This emits `if (!r) throwCAPIException()`
-    void checkAndThrowCAPIException(RewriterVar* r);
+    // This emits `if (r == exc_val) throwCAPIException()`
+    void checkAndThrowCAPIException(RewriterVar* r, int64_t exc_val = 0);
 
     void abort();
     void commit();

src/capi/typeobject.cpp

@@ -3247,7 +3247,7 @@ static Box* tppProxyToTpCall(Box* self, CallRewriteArgs* rewrite_args, ArgPassSp
         rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)self->cls->tp_call, rewrite_args->obj,
                                                              rewrite_args->arg1, rewrite_args->arg2);
         if (S == CXX)
-            rewrite_args->rewriter->call(true, (void*)checkAndThrowCAPIException);
+            rewrite_args->rewriter->checkAndThrowCAPIException(rewrite_args->out_rtn);
         rewrite_args->out_success = true;
     }
 

src/runtime/capi.cpp

@@ -797,6 +797,10 @@ void throwCAPIException() {
 }
 
 void checkAndThrowCAPIException() {
+    // Log these since these are expensive and usually avoidable:
+    static StatCounter num_checkAndThrowCAPIException("num_checkAndThrowCAPIException");
+    num_checkAndThrowCAPIException.log();
+
     Box* _type = cur_thread_state.curexc_type;
     if (!_type)
         assert(!cur_thread_state.curexc_value);
@@ -1572,7 +1576,7 @@ Box* BoxedCApiFunction::tppCall(Box* _self, CallRewriteArgs* rewrite_args, ArgPa
     }
 
     if (rewrite_args) {
-        rewrite_args->rewriter->call(false, (void*)checkAndThrowCAPIException);
+        rewrite_args->rewriter->checkAndThrowCAPIException(rewrite_args->out_rtn);
         rewrite_args->out_success = true;
     }
 

src/runtime/descr.cpp

@@ -387,7 +387,7 @@ Box* BoxedMethodDescriptor::tppCall(Box* _self, CallRewriteArgs* rewrite_args, A
         throwCAPIException();
 
     if (rewrite_args) {
-        rewrite_args->rewriter->call(false, (void*)checkAndThrowCAPIException);
+        rewrite_args->rewriter->checkAndThrowCAPIException(rewrite_args->out_rtn);
         rewrite_args->out_success = true;
     }
 
@@ -565,7 +565,7 @@ Box* BoxedWrapperObject::tppCall(Box* _self, CallRewriteArgs* rewrite_args, ArgP
             rewrite_args->out_rtn = rewriter->call(
                 true, (void*)wk, r_obj, rewrite_args->arg1,
                 rewriter->loadConst((intptr_t)self->descr->wrapped, Location::forArg(2)), rewrite_args->arg2);
-            rewriter->call(false, (void*)checkAndThrowCAPIException);
+            rewrite_args->rewriter->checkAndThrowCAPIException(rewrite_args->out_rtn);
             rewrite_args->out_success = true;
         }
     } else if (flags == PyWrapperFlag_PYSTON || flags == 0) {
@@ -577,7 +577,7 @@ Box* BoxedWrapperObject::tppCall(Box* _self, CallRewriteArgs* rewrite_args, ArgP
             rewrite_args->out_rtn
                 = rewriter->call(true, (void*)wrapper, r_obj, rewrite_args->arg1,
                                  rewriter->loadConst((intptr_t)self->descr->wrapped, Location::forArg(2)));
-            rewriter->call(false, (void*)checkAndThrowCAPIException);
+            rewrite_args->rewriter->checkAndThrowCAPIException(rewrite_args->out_rtn);
             rewrite_args->out_success = true;
         }
     } else {

src/runtime/objmodel.cpp

@@ -2056,12 +2056,11 @@ bool dataDescriptorSetSpecialCases(Box* obj, Box* val, Box* descr, SetattrRewrit
             args.push_back(r_obj);
             args.push_back(r_val);
             args.push_back(r_closure);
-            rewrite_args->rewriter->call(
+            RewriterVar* r_rtn = rewrite_args->rewriter->call(
                 /* has_side_effects */ true, (void*)getset_descr->set, args);
 
             if (descr->cls == capi_getset_cls)
-                // TODO I think we are supposed to check the return value?
-                rewrite_args->rewriter->call(true, (void*)checkAndThrowCAPIException);
+                rewrite_args->rewriter->checkAndThrowCAPIException(r_rtn, -1);
 
             rewrite_args->out_success = true;
         }
@@ -2578,7 +2577,11 @@ template <ExceptionStyle S> BoxedInt* lenInternal(Box* obj, LenRewriteArgs* rewr
             // support calling a RewriterVar (can only call fixed function addresses).
             r_m->addAttrGuard(offsetof(PySequenceMethods, sq_length), (intptr_t)m->sq_length);
             RewriterVar* r_n = rewrite_args->rewriter->call(true, (void*)m->sq_length, r_obj);
-            rewrite_args->rewriter->call(true, (void*)checkAndThrowCAPIException);
+
+            // Some CPython code seems to think that any negative return value means an exception,
+            // but the docs say -1. TODO it would be nice to just handle any negative value.
+            rewrite_args->rewriter->checkAndThrowCAPIException(r_n, -1);
+
             RewriterVar* r_r = rewrite_args->rewriter->call(false, (void*)boxInt, r_n);
 
             rewrite_args->out_success = true;
@@ -4231,7 +4234,8 @@ Box* compareInternal(Box* lhs, Box* rhs, int op_type, CompareRewriteArgs* rewrit
                     // support calling a RewriterVar (can only call fixed function addresses).
                     r_sqm->addAttrGuard(offsetof(PySequenceMethods, sq_contains), (intptr_t)sqm->sq_contains);
                     RewriterVar* r_b = rewrite_args->rewriter->call(true, (void*)sqm->sq_contains, r_rhs, r_lhs);
-                    rewrite_args->rewriter->call(true, (void*)checkAndThrowCAPIException);
+                    rewrite_args->rewriter->checkAndThrowCAPIException(r_b, -1);
+
                     // This could be inlined:
                     RewriterVar* r_r;
                     if (op_type == AST_TYPE::NotIn)