Nice, it's working

src/codegen/ast_interpreter.cpp

@@ -632,6 +632,8 @@ Value ASTInterpreter::visit_branch(AST_Branch* node) {
     ASSERT(v.o == True || v.o == False, "Should have called NONZERO before this branch");
 
     if (jit) {
+        jit->emitEndBlock();
+
         // Special note: emitSideExit decrefs v for us.
         // TODO: since the value is always True or False, maybe could optimize by putting the decref
         // before the conditional instead of after.
@@ -665,6 +667,7 @@ Value ASTInterpreter::visit_jump(AST_Jump* node) {
     if (jit) {
         if (backedge)
             jit->emitOSRPoint(node);
+        jit->emitEndBlock();
         jit->emitJump(node->target);
         finishJITing(node->target);
 
@@ -1036,6 +1039,7 @@ Value ASTInterpreter::visit_return(AST_Return* node) {
     Value s = node->value ? visit_expr(node->value) : getNone();
 
     if (jit) {
+        jit->emitEndBlock();
         jit->emitReturn(s);
         finishJITing();
     }
@@ -1738,8 +1742,10 @@ Box* ASTInterpreterJitInterface::derefHelper(void* _interpreter, InternedString
 Box* ASTInterpreterJitInterface::doOSRHelper(void* _interpreter, AST_Jump* node) {
     ASTInterpreter* interpreter = (ASTInterpreter*)_interpreter;
     ++interpreter->edgecount;
-    if (interpreter->edgecount >= OSR_THRESHOLD_BASELINE)
+    if (interpreter->edgecount >= OSR_THRESHOLD_BASELINE) {
+        // XXX refcounting here?
         return interpreter->doOSR(node);
+    }
     return NULL;
 }
 

src/codegen/baseline_jit.cpp

@@ -482,16 +482,6 @@ void JitFragmentWriter::emitExec(RewriterVar* code, RewriterVar* globals, Rewrit
 
 void JitFragmentWriter::emitJump(CFGBlock* b) {
     if (LOG_BJIT_ASSEMBLY) comment("BJIT: emitJump() start");
-    for (auto v : local_syms) {
-        if (v.second) {
-            if (LOG_BJIT_ASSEMBLY) {
-                // XXX silly but we need to keep this alive
-                std::string s = std::string("BJIT: decvref(") + v.first.c_str() + ")";
-                comment(*new std::string(s));
-            }
-            v.second->decvref(); // xdecref?
-        }
-    }
 
     RewriterVar* next = imm(b);
     addAction([=]() { _emitJump(b, next, exit_info); }, { next }, ActionType::NORMAL);
@@ -522,12 +512,20 @@ void JitFragmentWriter::emitRaise3(RewriterVar* arg0, RewriterVar* arg1, Rewrite
     call(false, (void*)raise3, arg0, arg1, arg2);
 }
 
-void JitFragmentWriter::emitReturn(RewriterVar* v) {
+void JitFragmentWriter::emitEndBlock() {
     for (auto v : local_syms) {
-        if (v.second)
-            v.second->decvref();
+        if (v.second) {
+            if (LOG_BJIT_ASSEMBLY) {
+                // XXX silly but we need to keep this alive
+                std::string s = std::string("BJIT: decvref(") + v.first.c_str() + ")";
+                comment(*new std::string(s));
+            }
+            v.second->decvref(); // xdecref?
         }
+    }
+}
 
+void JitFragmentWriter::emitReturn(RewriterVar* v) {
     v->stealRef();
     addAction([=]() { _emitReturn(v); }, { v }, ActionType::NORMAL);
     v->decvref();
@@ -603,10 +601,6 @@ void JitFragmentWriter::emitSideExit(STOLEN(RewriterVar*) v, Box* cmp_value, CFG
     vars.push_back(v);
     vars.push_back(var);
     vars.push_back(next_block_var);
-    for (auto v : local_syms) {
-        if (v.second)
-            vars.push_back(v.second);
-    }
 
     addAction([=]() { _emitSideExit(v, var, next_block, next_block_var); }, { v, var, next_block_var },
               ActionType::NORMAL);
@@ -701,10 +695,12 @@ int JitFragmentWriter::finishCompilation() {
         pp.release();
     }
 
+#ifndef NDEBUG
     if (LOG_BJIT_ASSEMBLY) {
         auto s = assembler->dump((uint8_t*)block->code/*, (uint8_t*)block->code + assembler->bytesWritten()*/);
         printf("%s\n", s.c_str());
     }
+#endif
 
     void* next_fragment_start = (uint8_t*)block->code + assembler->bytesWritten();
     if (exit_info.num_bytes)
@@ -1030,19 +1026,12 @@ void JitFragmentWriter::_emitSideExit(STOLEN(RewriterVar*) var, RewriterVar* val
         assembler->cmp(var_reg, assembler::Immediate(val));
     }
 
-    assert(0 && "the rewriter thinks this section is linear but it's not");
     {
         // TODO: Figure out if we need a large/small forward based on the number of local syms we will have to decref?
         assembler::LargeForwardJump jne(*assembler, assembler::COND_EQUAL);
 
-        //assert(0 && "hmm I think this is the issue");
         _decref(var);
 
-        for (auto v : local_syms) {
-            if (v.second)
-                _decref(v.second);
-        }
-
         ExitInfo exit_info;
         _emitJump(next_block, next_block_var, exit_info);
         if (exit_info.num_bytes) {
@@ -1055,7 +1044,6 @@ void JitFragmentWriter::_emitSideExit(STOLEN(RewriterVar*) var, RewriterVar* val
 
     if (LOG_BJIT_ASSEMBLY) assembler->comment("BJIT: decreffing emitSideExit var");
 
-    //assert(0 && "hmm I think this is the issue");
     _decref(var);
     if (LOG_BJIT_ASSEMBLY) assembler->comment("BJIT: decreffing emitSideExit var end");
 

src/codegen/baseline_jit.h

@@ -268,6 +268,7 @@ public:
     // emitSideExit steals a full ref from v, not just a vref
     void emitSideExit(STOLEN(RewriterVar*) v, Box* cmp_value, CFGBlock* next_block);
     void emitUncacheExcInfo();
+    void emitEndBlock();
 
     void abortCompilation();
     int finishCompilation();

src/core/options.cpp

@@ -25,7 +25,7 @@ int PYSTON_VERSION_MICRO = 0;
 int MAX_OPT_ITERATIONS = 1;
 
 bool LOG_IC_ASSEMBLY = false;
-bool LOG_BJIT_ASSEMBLY = true;
+bool LOG_BJIT_ASSEMBLY = false;
 
 bool CONTINUE_AFTER_FATAL = false;
 bool FORCE_INTERPRETER = true;