optimize list.__eq__ with a jit driver

fixes pypygh-5300

Author: cfbolz

pypy/module/pypyjit/test_pypy_c/test_containers.py

@@ -379,3 +379,28 @@ def main():
         opnames = log.opnames(ops)
         # only a call to ll_kvi, not to get_strategy_from_list_objects
         assert opnames.count("call_r") == 1
+
+    def test_list_eq(self):
+        def main():
+            l1 = [i * 3 + 1 for i in range(10000)]
+            l2 = [i * 3 + 1 for i in range(10000)]
+            return l1 == l2
+        log = self.run(main, [])
+        loop = log._filter(log.loops[-1], is_entry_bridge=False)
+        loop.match("""
+        i22 = uint_ge(i19, i6)
+        guard_false(i22, descr=...)
+        i23 = getarrayitem_gc_i(p8, i19, descr=...)
+        i24 = uint_ge(i19, i13)
+        guard_false(i24, descr=...)
+        i25 = getarrayitem_gc_i(p15, i19, descr=...)
+        i26 = int_eq(i25, i23)
+        guard_true(i26, descr=...)
+        i28 = int_add(i19, 1)
+        i29 = int_lt(i28, i6)
+        guard_true(i29, descr=...)
+        i30 = int_lt(i28, i13)
+        guard_true(i30, descr=...)
+        i31 = arraylen_gc(p8, descr=...)
+        i32 = arraylen_gc(p15, descr=...)
+        jump(i28, p1, p2, p5, i6, p8, p12, i13, p15, descr=...)""")

pypy/objspace/std/listobject.py

@@ -226,6 +226,68 @@ def list_unroll_condition(w_list1, space, w_list2):
     return (w_list1._unrolling_heuristic() or w_list2._unrolling_heuristic())
 
 
+def _get_printable_location(strategy_type1, strategy_type2, typ):
+    return 'list.eq [%s, %s, %s]' % (
+        strategy_type1,
+        strategy_type2,
+        typ)
+
+
+listeq_jitdriver = jit.JitDriver(
+    name='list.eq',
+    greens=['strategy_type1', 'strategy_type2', 'typ'],
+    reds='auto',
+    get_printable_location=_get_printable_location)
+
+def _make_list_eq(withjitdriver):
+    def list_eq(w_list, space, w_other):
+        # needs to be safe against eq_w() mutating the w_lists behind our back
+        length = w_list.length()
+        if length != w_other.length():
+            return space.w_False
+        if not length:
+            return space.w_True
+        if withjitdriver:
+            typ = type(w_list.getitem(0))
+
+        i = 0
+        while True:
+            if withjitdriver:
+                listeq_jitdriver.jit_merge_point(
+                        strategy_type1=type(w_list.strategy),
+                        strategy_type2=type(w_other.strategy),
+                        typ=typ,
+                        )
+            try:
+                w_item1 = w_list.getitem(i)
+                w_item2 = w_other.getitem(i)
+            except IndexError:
+                break
+            if not space.eq_w(w_item1, w_item2):
+                return space.w_False
+
+            i += 1
+
+        # if the list length is different now, the list was modified by eq_w with
+        l1 = w_list.length()
+        l2 = w_other.length()
+        if l1 != l2:
+            return space.w_False
+        return space.w_True
+    return list_eq
+_list_eq_withjitdriver = _make_list_eq(True)
+_list_eq_unroll = jit.unroll_safe(_make_list_eq(False))
+
+
+def list_eq(w_list, space, w_other):
+    # we can't use look_inside_iff because the jitdriver will be found in two
+    # different graphs then
+    if jit.we_are_jitted() and list_unroll_condition(w_list, space, w_other):
+        return _list_eq_unroll(w_list, space, w_other)
+    else:
+        return _list_eq_withjitdriver(w_list, space, w_other)
+
+
 class W_ListObject(W_Root):
     strategy = None
 
@@ -507,22 +569,7 @@ def descr_repr(self, space):
     def descr_eq(self, space, w_other):
         if not isinstance(w_other, W_ListObject):
             return space.w_NotImplemented
-        return self._descr_eq(space, w_other)
-
-    @jit.look_inside_iff(list_unroll_condition)
-    def _descr_eq(self, space, w_other):
-        # needs to be safe against eq_w() mutating the w_lists behind our back
-        if self.length() != w_other.length():
-            return space.w_False
-
-        # XXX in theory, this can be implemented more efficiently as well.
-        # let's not care for now
-        i = 0
-        while i < self.length() and i < w_other.length():
-            if not space.eq_w(self.getitem(i), w_other.getitem(i)):
-                return space.w_False
-            i += 1
-        return space.w_True
+        return list_eq(self, space, w_other)
 
     descr_ne = negate(descr_eq)
 

pypy/objspace/std/test/test_listobject.py

@@ -1544,6 +1544,30 @@ def __repr__(self):
         l.append(191)
         assert repr(l) == '[1, ouchie]'
 
+    def test_mutate_while_eq(self):
+        class Mean(object):
+            def __init__(self, i):
+                self.i = i
+            def __eq__(self, other):
+                if self.i == 9:
+                    del l1[self.i - 1]
+                return True
+        l1 = [Mean(i) for i in range(10)]
+        l2 = [Mean(i) for i in range(10)]
+        assert l1 != l2
+
+        class Mean(object):
+            def __init__(self, i):
+                self.i = i
+            def __eq__(self, other):
+                if self.i == 9:
+                    del l1[self.i - 1]
+                    del l2[self.i - 1]
+                return True
+        l1 = [Mean(i) for i in range(10)]
+        l2 = [Mean(i) for i in range(10)]
+        assert l1 == l2
+
     def test___getslice__(self):
         l = [1,2,3,4]
         res = l.__getslice__(0, 2)