misc fixes to get the bjit further

- fix instruction encoding of: add imm, mem
- fix Py_REF_DEBUG when increfing multiple times at once
- add missing bjit annotations

src/asm_writing/assembler.cpp

@@ -116,26 +116,28 @@ void Assembler::emitArith(Immediate imm, Indirect mem, int opcode) {
         mem_idx -= 8;
     }
 
-    // TODO: needs testing, then remove this trap
-    trap();
-
     emitRex(rex);
 
     bool needssib = (mem_idx == 0b100);
     assert(!needssib && "untested");
     int mode = getModeFromOffset(mem.offset, mem_idx);
+    assert(mode != 0b10 && "not yet supported");
 
     if (-0x80 <= amount && amount < 0x80) {
         emitByte(0x83);
         if (needssib)
             emitSIB(0b00, 0b100, mem_idx);
         emitModRM(mode, opcode, mem_idx);
+        if (mode == 0b01)
+            emitByte(mem.offset);
         emitByte(amount);
     } else {
         emitByte(0x81);
         if (needssib)
             emitSIB(0b00, 0b100, mem_idx);
         emitModRM(mode, opcode, mem_idx);
+        if (mode == 0b01)
+            emitByte(mem.offset);
         emitInt(amount, 4);
     }
 }

src/asm_writing/rewriter.cpp

@@ -542,6 +542,7 @@ void Rewriter::_incref(RewriterVar* var, int num_refs) {
 // llvm::ArrayRef<RewriterVar*>());
 #ifdef Py_REF_DEBUG
     // assembler->trap();
+    for (int i = 0; i < num_refs; ++i)
         assembler->incq(assembler::Immediate(&_Py_RefTotal));
 #endif
     auto reg = var->getInReg();

src/codegen/ast_interpreter.cpp

@@ -1798,7 +1798,8 @@ void ASTInterpreterJitInterface::pendingCallsCheckHelper() {
 #endif
 }
 
-void ASTInterpreterJitInterface::setExcInfoHelper(void* _interpreter, Box* type, Box* value, Box* traceback) {
+void ASTInterpreterJitInterface::setExcInfoHelper(void* _interpreter, STOLEN(Box*) type, STOLEN(Box*) value,
+                                                  STOLEN(Box*) traceback) {
     ASTInterpreter* interpreter = (ASTInterpreter*)_interpreter;
     setFrameExcInfo(interpreter->getFrameInfo(), type, value, traceback);
 }

src/codegen/ast_interpreter.h

@@ -49,7 +49,7 @@ struct ASTInterpreterJitInterface {
     static Box* derefHelper(void* interp, InternedString s);
     static Box* landingpadHelper(void* interp);
     static void pendingCallsCheckHelper();
-    static void setExcInfoHelper(void* interp, Box* type, Box* value, Box* traceback);
+    static void setExcInfoHelper(void* interp, STOLEN(Box*) type, STOLEN(Box*) value, STOLEN(Box*) traceback);
     static void setLocalClosureHelper(void* interp, long vreg, InternedString id, Box* v);
     static void uncacheExcInfoHelper(void* interp);
     static void raise0Helper(void* interp) __attribute__((noreturn));

src/codegen/baseline_jit.cpp

@@ -539,6 +539,9 @@ void JitFragmentWriter::emitRaise0() {
 
 void JitFragmentWriter::emitRaise3(RewriterVar* arg0, RewriterVar* arg1, RewriterVar* arg2) {
     call(false, (void*)raise3, arg0, arg1, arg2);
+    arg0->refConsumed();
+    arg1->refConsumed();
+    arg2->refConsumed();
 }
 
 void JitFragmentWriter::emitEndBlock() {
@@ -570,6 +573,9 @@ void JitFragmentWriter::emitSetCurrentInst(AST_stmt* node) {
 
 void JitFragmentWriter::emitSetExcInfo(RewriterVar* type, RewriterVar* value, RewriterVar* traceback) {
     call(false, (void*)ASTInterpreterJitInterface::setExcInfoHelper, getInterp(), type, value, traceback);
+    type->refConsumed();
+    value->refConsumed();
+    traceback->refConsumed();
 }
 
 void JitFragmentWriter::emitSetGlobal(Box* global, BoxedString* s, STOLEN(RewriterVar*) v) {
@@ -887,7 +893,6 @@ Box* JitFragmentWriter::runtimeCallHelper(Box* obj, ArgPassSpec argspec, TypeRec
 void JitFragmentWriter::_emitGetLocal(RewriterVar* val_var, const char* name) {
     assembler::Register var_reg = val_var->getInReg();
     assembler->test(var_reg, var_reg);
-    val_var->bumpUse();
 
     {
         assembler::ForwardJump jnz(*assembler, assembler::COND_NOT_ZERO);
@@ -897,6 +902,7 @@ void JitFragmentWriter::_emitGetLocal(RewriterVar* val_var, const char* name) {
     }
 
     _incref(val_var);
+    val_var->bumpUse();
 }
 
 void JitFragmentWriter::_emitJump(CFGBlock* b, RewriterVar* block_next, ExitInfo& exit_info) {