pyston · kmod · Aug 10, 2015 · Aug 10, 2015
diff --git a/src/codegen/patchpoints.cpp b/src/codegen/patchpoints.cpp
@@ -333,7 +333,7 @@ ICSetupInfo* createBinexpIC(TypeRecorder* type_recorder) {
 }
 
 ICSetupInfo* createNonzeroIC(TypeRecorder* type_recorder) {
-    return ICSetupInfo::initialize(true, 2, 64, ICSetupInfo::Nonzero, type_recorder);
+    return ICSetupInfo::initialize(true, 2, 512, ICSetupInfo::Nonzero, type_recorder);
 }
 
 ICSetupInfo* createHasnextIC(TypeRecorder* type_recorder) {

diff --git a/src/runtime/ics.h b/src/runtime/ics.h
@@ -86,7 +86,7 @@ class BinopIC : public RuntimeIC {
 
 class NonzeroIC : public RuntimeIC {
 public:
-    NonzeroIC() : RuntimeIC((void*)nonzero, 1, 40) {}
+    NonzeroIC() : RuntimeIC((void*)nonzero, 1, 512) {}
 
     bool call(Box* obj) { return call_bool(obj); }
 };

diff --git a/src/runtime/objmodel.cpp b/src/runtime/objmodel.cpp
@@ -2275,12 +2275,28 @@ extern "C" void setattr(Box* obj, BoxedString* attr, Box* attr_val) {
     }
 }
 
+static bool nonzeroHelper(Box* r) {
+    // I believe this behavior is handled by the slot wrappers in CPython:
+    if (r->cls == bool_cls) {
+        BoxedBool* b = static_cast<BoxedBool*>(r);
+        bool rtn = b->n;
+        return rtn;
+    } else if (r->cls == int_cls) {
+        BoxedInt* b = static_cast<BoxedInt*>(r);
+        bool rtn = b->n != 0;
+        return rtn;
+    } else {
+        raiseExcHelper(TypeError, "__nonzero__ should return bool or int, returned %s", getTypeName(r));
+    }
+}
+
 extern "C" bool nonzero(Box* obj) {
     STAT_TIMER(t0, "us_timer_slowpath_nonzero", 10);
 
     assert(gc::isValidGCObject(obj));
 
     static StatCounter slowpath_nonzero("slowpath_nonzero");
+    slowpath_nonzero.log();
 
     std::unique_ptr<Rewriter> rewriter(
         Rewriter::createRewriter(__builtin_extract_return_addr(__builtin_return_address(0)), 1, "nonzero"));
@@ -2374,40 +2390,47 @@ extern "C" bool nonzero(Box* obj) {
         return r;
     }
 
-    // TODO: rewrite these.
     static BoxedString* nonzero_str = internStringImmortal("__nonzero__");
     static BoxedString* len_str = internStringImmortal("__len__");
-    // go through descriptor logic
-    Box* func = getclsattrInternal(obj, nonzero_str, NULL);
-    if (!func)
-        func = getclsattrInternal(obj, len_str, NULL);
-
-    if (func == NULL) {
-        ASSERT(obj->cls->is_user_defined || obj->cls == classobj_cls || obj->cls == type_cls
-                   || isSubclass(obj->cls, Exception) || obj->cls == file_cls || obj->cls == traceback_cls
-                   || obj->cls == instancemethod_cls || obj->cls == module_cls || obj->cls == capifunc_cls
-                   || obj->cls == builtin_function_or_method_cls || obj->cls == method_cls || obj->cls == frame_cls
-                   || obj->cls == generator_cls || obj->cls == capi_getset_cls || obj->cls == pyston_getset_cls
-                   || obj->cls == wrapperdescr_cls,
-               "%s.__nonzero__", getTypeName(obj)); // TODO
-
-        // TODO should rewrite these?
-        return true;
-    }
 
-    Box* r = runtimeCallInternal<CXX>(func, NULL, ArgPassSpec(0), NULL, NULL, NULL, NULL, NULL);
-    // I believe this behavior is handled by the slot wrappers in CPython:
-    if (r->cls == bool_cls) {
-        BoxedBool* b = static_cast<BoxedBool*>(r);
-        bool rtn = b->n;
-        return rtn;
-    } else if (r->cls == int_cls) {
-        BoxedInt* b = static_cast<BoxedInt*>(r);
-        bool rtn = b->n != 0;
-        return rtn;
-    } else {
-        raiseExcHelper(TypeError, "__nonzero__ should return bool or int, returned %s", getTypeName(r));
+    // try __nonzero__
+    GetattrRewriteArgs grewrite_args(rewriter.get(), r_obj, rewriter ? rewriter->getReturnDestination() : Location());
+    Box* func = getclsattrInternal(obj, nonzero_str, rewriter ? &grewrite_args : NULL);
+    if (!grewrite_args.out_success)
+        rewriter.reset();
+
+    if (!func) {
+        // try __len__
+        grewrite_args
+            = GetattrRewriteArgs(rewriter.get(), r_obj, rewriter ? rewriter->getReturnDestination() : Location());
+        func = getclsattrInternal(obj, len_str, rewriter ? &grewrite_args : NULL);
+        if (!grewrite_args.out_success)
+            rewriter.reset();
+
+        if (func == NULL) {
+            ASSERT(obj->cls->is_user_defined || obj->cls == classobj_cls || obj->cls == type_cls
+                       || isSubclass(obj->cls, Exception) || obj->cls == file_cls || obj->cls == traceback_cls
+                       || obj->cls == instancemethod_cls || obj->cls == module_cls || obj->cls == capifunc_cls
+                       || obj->cls == builtin_function_or_method_cls || obj->cls == method_cls || obj->cls == frame_cls
+                       || obj->cls == generator_cls || obj->cls == capi_getset_cls || obj->cls == pyston_getset_cls
+                       || obj->cls == wrapperdescr_cls,
+                   "%s.__nonzero__", getTypeName(obj)); // TODO
+
+            if (rewriter.get()) {
+                RewriterVar* b = rewriter->loadConst(1, rewriter->getReturnDestination());
+                rewriter->commitReturning(b);
+            }
+            return true;
+        }
+    }
+    CallRewriteArgs cargs(rewriter.get(), grewrite_args.out_rtn,
+                          rewriter ? rewriter->getReturnDestination() : Location());
+    Box* rtn = runtimeCallInternal0<CXX>(func, rewriter ? &cargs : NULL, ArgPassSpec(0));
+    if (cargs.out_success) {
+        RewriterVar* b = rewriter->call(false, (void*)nonzeroHelper, cargs.out_rtn);
+        rewriter->commitReturning(b);
     }
+    return nonzeroHelper(rtn);
 }
 
 extern "C" BoxedString* str(Box* obj) {

diff --git a/test/tests/nonzero_patching.py b/test/tests/nonzero_patching.py
@@ -1,5 +1,5 @@
 # run_args: -n
-# statcheck: noninit_count('slowpath_nonzero') <= 10
+# statcheck: noninit_count('slowpath_nonzero') <= 25
 
 def f():
     for i in xrange(-10, 10):
@@ -9,3 +9,29 @@ def f():
         else:
             print "is false-y"
 f()
+
+num_nonzero = 0
+num_len = 0
+class C(object):
+	def mynonzero(self):
+		global num_nonzero
+		num_nonzero += 1
+		return True
+	def mylen(self):
+		global num_len
+		num_len += 1
+		return True
+f = C()
+s = 0
+for i in range(1000):
+	if f:
+		pass
+	if i == 200:
+		C.__len__ = C.mylen
+	if i == 400:
+		C.__nonzero__ = C.mynonzero
+	if i == 600:
+		del C.__len__
+	if i == 800:
+		del C.__nonzero__
+print num_nonzero, num_len