bjit: better hasnext handling

keep track of which variables are known to be a python bool so that we can remove some of the nonzero() calls

src/codegen/baseline_jit.cpp

@@ -438,7 +438,9 @@ RewriterVar* JitFragmentWriter::emitGetPystonIter(RewriterVar* v) {
 }
 
 RewriterVar* JitFragmentWriter::emitHasnext(RewriterVar* v) {
-    return call(false, (void*)hasnextHelper, v)->setType(RefType::OWNED);
+    auto rtn = call(false, (void*)hasnextHelper, v)->setType(RefType::BORROWED);
+    var_is_a_python_bool.insert(rtn);
+    return rtn;
 }
 
 RewriterVar* JitFragmentWriter::emitImportFrom(RewriterVar* module, BoxedString* name) {
@@ -460,12 +462,18 @@ RewriterVar* JitFragmentWriter::emitLandingpad() {
 }
 
 RewriterVar* JitFragmentWriter::emitNonzero(RewriterVar* v) {
+    if (var_is_a_python_bool.count(v))
+        return v;
     // nonzeroHelper returns bool
-    return call(false, (void*)nonzeroHelper, v)->setType(RefType::BORROWED);
+    auto rtn = call(false, (void*)nonzeroHelper, v)->setType(RefType::BORROWED);
+    var_is_a_python_bool.insert(rtn);
+    return rtn;
 }
 
 RewriterVar* JitFragmentWriter::emitNotNonzero(RewriterVar* v) {
-    return call(false, (void*)notHelper, v)->setType(RefType::BORROWED);
+    auto rtn = call(false, (void*)notHelper, v)->setType(RefType::BORROWED);
+    var_is_a_python_bool.insert(rtn);
+    return rtn;
 }
 
 RewriterVar* JitFragmentWriter::emitRepr(RewriterVar* v) {
@@ -1002,8 +1010,8 @@ Box* JitFragmentWriter::exceptionMatchesHelper(Box* obj, Box* cls) {
     return boxBool(exceptionMatches(obj, cls));
 }
 
-Box* JitFragmentWriter::hasnextHelper(Box* b) {
-    return boxBool(pyston::hasnext(b));
+BORROWED(Box*) JitFragmentWriter::hasnextHelper(Box* b) {
+    return pyston::hasnext(b) ? Py_True : Py_False;
 }
 
 BORROWED(Box*) JitFragmentWriter::nonzeroHelper(Box* b) {
@@ -1208,7 +1216,6 @@ void JitFragmentWriter::_emitSideExit(STOLEN(RewriterVar*) var, RewriterVar* val
         // Hax: override the automatic refcount system
         var->reftype = RefType::BORROWED;
     }
-    assert(var->num_refs_consumed == 0);
 
     assembler::Register var_reg = var->getInReg();
     if (isLargeConstant(val)) {

src/codegen/baseline_jit.h

@@ -211,6 +211,7 @@ private:
     // TODO: in the future we could reuse this information between different basic blocks
     llvm::DenseSet<int> known_non_null_vregs;
     std::unique_ptr<ICInfo> ic_info;
+    llvm::SmallPtrSet<RewriterVar*, 4> var_is_a_python_bool;
 
     // Optional points to a CFGBlock and a patch location which should get patched to a direct jump if
     // the specified block gets JITed. The patch location is guaranteed to be at least 'min_patch_size' bytes long.
@@ -333,7 +334,7 @@ private:
     static Box* createSetHelper(uint64_t num, Box** data);
     static Box* createTupleHelper(uint64_t num, Box** data);
     static Box* exceptionMatchesHelper(Box* obj, Box* cls);
-    static Box* hasnextHelper(Box* b);
+    static BORROWED(Box*) hasnextHelper(Box* b);
     static BORROWED(Box*) nonzeroHelper(Box* b);
     static BORROWED(Box*) notHelper(Box* b);
     static Box* runtimeCallHelper(Box* obj, ArgPassSpec argspec, TypeRecorder* type_recorder, Box** args,