[mypyc] Speed up implicit __ne__ method

mypyc/codegen/emitfunc.py

@@ -89,7 +89,7 @@
     RStruct,
     RTuple,
     RType,
-    is_bool_rprimitive,
+    is_bool_or_bit_rprimitive,
     is_int32_rprimitive,
     is_int64_rprimitive,
     is_int_rprimitive,
@@ -633,7 +633,7 @@ def emit_method_call(self, dest: str, op_obj: Value, name: str, op_args: list[Va
     def visit_inc_ref(self, op: IncRef) -> None:
         if (
             isinstance(op.src, Box)
-            and (is_none_rprimitive(op.src.src.type) or is_bool_rprimitive(op.src.src.type))
+            and (is_none_rprimitive(op.src.src.type) or is_bool_or_bit_rprimitive(op.src.src.type))
             and HAVE_IMMORTAL
         ):
             # On Python 3.12+, None/True/False are immortal, and we can skip inc ref

mypyc/irbuild/classdef.py

@@ -845,7 +845,9 @@ def gen_glue_ne_method(builder: IRBuilder, cls: ClassIR, line: int) -> None:
             )
             builder.activate_block(regular_block)
             rettype = bool_rprimitive if return_bool and strict_typing else object_rprimitive
-            retval = builder.coerce(builder.unary_op(eqval, "not", line), rettype, line)
+            retval = builder.coerce(
+                builder.builder.unary_not(eqval, line, likely_bool=True), rettype, line
+            )
             builder.add(Return(retval))
             builder.activate_block(not_implemented_block)
             builder.add(Return(not_implemented))

mypyc/irbuild/ll_builder.py

@@ -56,6 +56,7 @@
     KeepAlive,
     LoadAddress,
     LoadErrorValue,
+    LoadGlobal,
     LoadLiteral,
     LoadMem,
     LoadStatic,
@@ -108,6 +109,7 @@
     is_int_rprimitive,
     is_list_rprimitive,
     is_none_rprimitive,
+    is_object_rprimitive,
     is_set_rprimitive,
     is_short_int_rprimitive,
     is_str_rprimitive,
@@ -1318,6 +1320,14 @@ def none_object(self) -> Value:
         """Load Python None value (type: object_rprimitive)."""
         return self.add(LoadAddress(none_object_op.type, none_object_op.src, line=-1))
 
+    def true_object(self) -> Value:
+        """Load Python True object (type: object_rprimitive)."""
+        return self.add(LoadGlobal(object_rprimitive, "Py_True"))
+
+    def false_object(self) -> Value:
+        """Load Python False object (type: object_rprimitive)."""
+        return self.add(LoadGlobal(object_rprimitive, "Py_False"))
+
     def load_int(self, value: int) -> Value:
         """Load a tagged (Python) integer literal value."""
         if value > MAX_LITERAL_SHORT_INT or value < MIN_LITERAL_SHORT_INT:
@@ -1663,12 +1673,39 @@ def _non_specialized_unary_op(self, value: Value, op: str, line: int) -> Value:
         assert target, "Unsupported unary operation: %s" % op
         return target
 
-    def unary_not(self, value: Value, line: int) -> Value:
-        """Perform unary 'not'."""
+    def unary_not(self, value: Value, line: int, *, likely_bool: bool = False) -> Value:
+        """Perform unary 'not'.
+
+        Args:
+            likely_bool: The operand is likely a bool value, even if the type is something
+                more general, so specialize for bool values
+        """
         typ = value.type
         if is_bool_or_bit_rprimitive(typ):
             mask = Integer(1, typ, line)
             return self.int_op(typ, value, mask, IntOp.XOR, line)
+        if likely_bool and is_object_rprimitive(typ):
+            # First quickly check if it's a bool, and otherwise fall back to generic op.
+            res = Register(bit_rprimitive)
+            false, not_false, true, other = BasicBlock(), BasicBlock(), BasicBlock(), BasicBlock()
+            out = BasicBlock()
+            cmp = self.add(ComparisonOp(value, self.true_object(), ComparisonOp.EQ, line))
+            self.add(Branch(cmp, false, not_false, Branch.BOOL))
+            self.activate_block(false)
+            self.add(Assign(res, self.false()))
+            self.goto(out)
+            self.activate_block(not_false)
+            cmp = self.add(ComparisonOp(value, self.false_object(), ComparisonOp.EQ, line))
+            self.add(Branch(cmp, true, other, Branch.BOOL))
+            self.activate_block(true)
+            self.add(Assign(res, self.true()))
+            self.goto(out)
+            self.activate_block(other)
+            val = self._non_specialized_unary_op(value, "not", line)
+            self.add(Assign(res, val))
+            self.goto(out)
+            self.activate_block(out)
+            return res
         return self._non_specialized_unary_op(value, "not", line)
 
     def unary_minus(self, value: Value, line: int) -> Value:

mypyc/test-data/irbuild-basic.test

@@ -2347,22 +2347,42 @@ def A.__ne__(__mypyc_self__, rhs):
     __mypyc_self__ :: __main__.A
     rhs, r0, r1 :: object
     r2 :: bit
-    r3 :: i32
-    r4 :: bit
-    r5 :: bool
+    r3 :: object
+    r4, r5 :: bit
     r6 :: object
+    r7 :: bit
+    r8 :: i32
+    r9 :: bit
+    r10 :: bool
+    r11 :: object
 L0:
     r0 = __mypyc_self__.__eq__(rhs)
     r1 = load_address _Py_NotImplementedStruct
     r2 = r0 == r1
-    if r2 goto L2 else goto L1 :: bool
+    if r2 goto L7 else goto L1 :: bool
 L1:
-    r3 = PyObject_Not(r0)
-    r4 = r3 >= 0 :: signed
-    r5 = truncate r3: i32 to builtins.bool
-    r6 = box(bool, r5)
-    return r6
+    r3 = load_global Py_True :: static
+    r4 = r0 == r3
+    if r4 goto L2 else goto L3 :: bool
 L2:
+    r5 = 0
+    goto L6
+L3:
+    r6 = load_global Py_False :: static
+    r7 = r0 == r6
+    if r7 goto L4 else goto L5 :: bool
+L4:
+    r5 = 1
+    goto L6
+L5:
+    r8 = PyObject_Not(r0)
+    r9 = r8 >= 0 :: signed
+    r10 = truncate r8: i32 to builtins.bool
+    r5 = r10
+L6:
+    r11 = box(bit, r5)
+    return r11
+L7:
     return r1
 
 [case testDecorators_toplevel]

mypyc/test-data/irbuild-classes.test

@@ -854,22 +854,42 @@ def Base.__ne__(__mypyc_self__, rhs):
     __mypyc_self__ :: __main__.Base
     rhs, r0, r1 :: object
     r2 :: bit
-    r3 :: i32
-    r4 :: bit
-    r5 :: bool
+    r3 :: object
+    r4, r5 :: bit
     r6 :: object
+    r7 :: bit
+    r8 :: i32
+    r9 :: bit
+    r10 :: bool
+    r11 :: object
 L0:
     r0 = __mypyc_self__.__eq__(rhs)
     r1 = load_address _Py_NotImplementedStruct
     r2 = r0 == r1
-    if r2 goto L2 else goto L1 :: bool
+    if r2 goto L7 else goto L1 :: bool
 L1:
-    r3 = PyObject_Not(r0)
-    r4 = r3 >= 0 :: signed
-    r5 = truncate r3: i32 to builtins.bool
-    r6 = box(bool, r5)
-    return r6
+    r3 = load_global Py_True :: static
+    r4 = r0 == r3
+    if r4 goto L2 else goto L3 :: bool
 L2:
+    r5 = 0
+    goto L6
+L3:
+    r6 = load_global Py_False :: static
+    r7 = r0 == r6
+    if r7 goto L4 else goto L5 :: bool
+L4:
+    r5 = 1
+    goto L6
+L5:
+    r8 = PyObject_Not(r0)
+    r9 = r8 >= 0 :: signed
+    r10 = truncate r8: i32 to builtins.bool
+    r5 = r10
+L6:
+    r11 = box(bit, r5)
+    return r11
+L7:
     return r1
 def Derived.__eq__(self, other):
     self :: __main__.Derived
@@ -979,22 +999,42 @@ def Derived.__ne__(__mypyc_self__, rhs):
     __mypyc_self__ :: __main__.Derived
     rhs, r0, r1 :: object
     r2 :: bit
-    r3 :: i32
-    r4 :: bit
-    r5 :: bool
+    r3 :: object
+    r4, r5 :: bit
     r6 :: object
+    r7 :: bit
+    r8 :: i32
+    r9 :: bit
+    r10 :: bool
+    r11 :: object
 L0:
     r0 = __mypyc_self__.__eq__(rhs)
     r1 = load_address _Py_NotImplementedStruct
     r2 = r0 == r1
-    if r2 goto L2 else goto L1 :: bool
+    if r2 goto L7 else goto L1 :: bool
 L1:
-    r3 = PyObject_Not(r0)
-    r4 = r3 >= 0 :: signed
-    r5 = truncate r3: i32 to builtins.bool
-    r6 = box(bool, r5)
-    return r6
+    r3 = load_global Py_True :: static
+    r4 = r0 == r3
+    if r4 goto L2 else goto L3 :: bool
 L2:
+    r5 = 0
+    goto L6
+L3:
+    r6 = load_global Py_False :: static
+    r7 = r0 == r6
+    if r7 goto L4 else goto L5 :: bool
+L4:
+    r5 = 1
+    goto L6
+L5:
+    r8 = PyObject_Not(r0)
+    r9 = r8 >= 0 :: signed
+    r10 = truncate r8: i32 to builtins.bool
+    r5 = r10
+L6:
+    r11 = box(bit, r5)
+    return r11
+L7:
     return r1
 
 [case testDefaultVars]

mypyc/test-data/run-dunders-special.test

@@ -8,3 +8,5 @@ class UsesNotImplemented:
 
 def test_not_implemented() -> None:
     assert UsesNotImplemented() != object()
+    x = UsesNotImplemented() == object()
+    assert not x

mypyc/test-data/run-dunders.test

@@ -965,3 +965,25 @@ def test_final() -> None:
     assert b + 3 == 9
     assert (a < A(5)) is False
     assert (b < A(5)) is True
+
+[case testDundersEq]
+class Eq:
+    def __init__(self, x: int) -> None:
+        self.x = x
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Eq):
+            return NotImplemented
+        return self.x == other.x
+
+def eq(x: Eq, y: Eq) -> bool:
+    return x == y
+
+def ne(x: Eq, y: Eq) -> bool:
+    return x != y
+
+def test_equality_with_implicit_ne() -> None:
+    assert eq(Eq(1), Eq(1))
+    assert not eq(Eq(1), Eq(2))
+    assert ne(Eq(1), Eq(2))
+    assert not ne(Eq(1), Eq(1))