[mypyc] Add is_bool_or_bit_rprimitive (#19406)

Added a wrapper to check if a type is either a bool or bit primitive as
these two checks are often done together.

The wrapper should help in preventing suboptimal code generation if one
forgets to check for the bit primitive in cases when it can be trivially
expanded to bool. One such case was in translation of binary ops, which
is fixed in this PR.

Example code:
```
 def f(a: float, b: float, c: float) -> bool:
     return (a == b) & (a == c)
```

IR before:
```
def f(a, b, c):
    a, b, c :: float
    r0, r1 :: bit
    r2 :: bool
    r3 :: int
    r4 :: bool
    r5, r6 :: int
    r7 :: object
    r8, r9 :: bool
L0:
    r0 = a == b
    r1 = a == c
    r2 = r0 << 1
    r3 = extend r2: builtins.bool to builtins.int
    r4 = r1 << 1
    r5 = extend r4: builtins.bool to builtins.int
    r6 = CPyTagged_And(r3, r5)
    dec_ref r3 :: int
    dec_ref r5 :: int
    r7 = box(int, r6)
    r8 = unbox(bool, r7)
    dec_ref r7
    if is_error(r8) goto L2 (error at f:2) else goto L1
L1:
    return r8
L2:
    r9 = <error> :: bool
    return r9
```

IR after:
```
def f(a, b, c):
    a, b, c :: float
    r0, r1 :: bit
    r2 :: bool
L0:
    r0 = a == b
    r1 = a == c
    r2 = r0 & r1
    return r2
```

Author: p-sawicki

mypyc/codegen/emit.py

@@ -28,8 +28,7 @@
     RType,
     RUnion,
     int_rprimitive,
-    is_bit_rprimitive,
-    is_bool_rprimitive,
+    is_bool_or_bit_rprimitive,
     is_bytes_rprimitive,
     is_dict_rprimitive,
     is_fixed_width_rtype,
@@ -615,8 +614,7 @@ def emit_cast(
             or is_range_rprimitive(typ)
             or is_float_rprimitive(typ)
             or is_int_rprimitive(typ)
-            or is_bool_rprimitive(typ)
-            or is_bit_rprimitive(typ)
+            or is_bool_or_bit_rprimitive(typ)
             or is_fixed_width_rtype(typ)
         ):
             if declare_dest:
@@ -638,7 +636,7 @@ def emit_cast(
             elif is_int_rprimitive(typ) or is_fixed_width_rtype(typ):
                 # TODO: Range check for fixed-width types?
                 prefix = "PyLong"
-            elif is_bool_rprimitive(typ) or is_bit_rprimitive(typ):
+            elif is_bool_or_bit_rprimitive(typ):
                 prefix = "PyBool"
             else:
                 assert False, f"unexpected primitive type: {typ}"
@@ -889,7 +887,7 @@ def emit_unbox(
             self.emit_line("else {")
             self.emit_line(failure)
             self.emit_line("}")
-        elif is_bool_rprimitive(typ) or is_bit_rprimitive(typ):
+        elif is_bool_or_bit_rprimitive(typ):
             # Whether we are borrowing or not makes no difference.
             if declare_dest:
                 self.emit_line(f"char {dest};")
@@ -1015,7 +1013,7 @@ def emit_box(
         if is_int_rprimitive(typ) or is_short_int_rprimitive(typ):
             # Steal the existing reference if it exists.
             self.emit_line(f"{declaration}{dest} = CPyTagged_StealAsObject({src});")
-        elif is_bool_rprimitive(typ) or is_bit_rprimitive(typ):
+        elif is_bool_or_bit_rprimitive(typ):
             # N.B: bool is special cased to produce a borrowed value
             # after boxing, so we don't need to increment the refcount
             # when this comes directly from a Box op.

mypyc/ir/ops.py

@@ -42,8 +42,7 @@ class to enable the new behavior. Sometimes adding a new abstract
     cstring_rprimitive,
     float_rprimitive,
     int_rprimitive,
-    is_bit_rprimitive,
-    is_bool_rprimitive,
+    is_bool_or_bit_rprimitive,
     is_int_rprimitive,
     is_none_rprimitive,
     is_pointer_rprimitive,
@@ -1089,11 +1088,7 @@ def __init__(self, src: Value, line: int = -1) -> None:
         self.src = src
         self.type = object_rprimitive
         # When we box None and bool values, we produce a borrowed result
-        if (
-            is_none_rprimitive(self.src.type)
-            or is_bool_rprimitive(self.src.type)
-            or is_bit_rprimitive(self.src.type)
-        ):
+        if is_none_rprimitive(self.src.type) or is_bool_or_bit_rprimitive(self.src.type):
             self.is_borrowed = True
 
     def sources(self) -> list[Value]:

mypyc/ir/rtypes.py

@@ -582,6 +582,10 @@ def is_bit_rprimitive(rtype: RType) -> bool:
     return isinstance(rtype, RPrimitive) and rtype.name == "bit"
 
 
+def is_bool_or_bit_rprimitive(rtype: RType) -> bool:
+    return is_bool_rprimitive(rtype) or is_bit_rprimitive(rtype)
+
+
 def is_object_rprimitive(rtype: RType) -> bool:
     return isinstance(rtype, RPrimitive) and rtype.name == "builtins.object"
 

mypyc/irbuild/ll_builder.py

@@ -93,8 +93,7 @@
     dict_rprimitive,
     float_rprimitive,
     int_rprimitive,
-    is_bit_rprimitive,
-    is_bool_rprimitive,
+    is_bool_or_bit_rprimitive,
     is_bytes_rprimitive,
     is_dict_rprimitive,
     is_fixed_width_rtype,
@@ -381,16 +380,12 @@ def coerce(
             ):
                 # Equivalent types
                 return src
-            elif (is_bool_rprimitive(src_type) or is_bit_rprimitive(src_type)) and is_tagged(
-                target_type
-            ):
+            elif is_bool_or_bit_rprimitive(src_type) and is_tagged(target_type):
                 shifted = self.int_op(
                     bool_rprimitive, src, Integer(1, bool_rprimitive), IntOp.LEFT_SHIFT
                 )
                 return self.add(Extend(shifted, target_type, signed=False))
-            elif (
-                is_bool_rprimitive(src_type) or is_bit_rprimitive(src_type)
-            ) and is_fixed_width_rtype(target_type):
+            elif is_bool_or_bit_rprimitive(src_type) and is_fixed_width_rtype(target_type):
                 return self.add(Extend(src, target_type, signed=False))
             elif isinstance(src, Integer) and is_float_rprimitive(target_type):
                 if is_tagged(src_type):
@@ -1341,7 +1336,11 @@ def binary_op(self, lreg: Value, rreg: Value, op: str, line: int) -> Value:
             return self.compare_strings(lreg, rreg, op, line)
         if is_bytes_rprimitive(ltype) and is_bytes_rprimitive(rtype) and op in ("==", "!="):
             return self.compare_bytes(lreg, rreg, op, line)
-        if is_bool_rprimitive(ltype) and is_bool_rprimitive(rtype) and op in BOOL_BINARY_OPS:
+        if (
+            is_bool_or_bit_rprimitive(ltype)
+            and is_bool_or_bit_rprimitive(rtype)
+            and op in BOOL_BINARY_OPS
+        ):
             if op in ComparisonOp.signed_ops:
                 return self.bool_comparison_op(lreg, rreg, op, line)
             else:
@@ -1355,7 +1354,7 @@ def binary_op(self, lreg: Value, rreg: Value, op: str, line: int) -> Value:
                     op_id = int_op_to_id[op]
                 else:
                     op_id = IntOp.DIV
-                if is_bool_rprimitive(rtype) or is_bit_rprimitive(rtype):
+                if is_bool_or_bit_rprimitive(rtype):
                     rreg = self.coerce(rreg, ltype, line)
                     rtype = ltype
                 if is_fixed_width_rtype(rtype) or is_tagged(rtype):
@@ -1367,7 +1366,7 @@ def binary_op(self, lreg: Value, rreg: Value, op: str, line: int) -> Value:
             elif op in ComparisonOp.signed_ops:
                 if is_int_rprimitive(rtype):
                     rreg = self.coerce_int_to_fixed_width(rreg, ltype, line)
-                elif is_bool_rprimitive(rtype) or is_bit_rprimitive(rtype):
+                elif is_bool_or_bit_rprimitive(rtype):
                     rreg = self.coerce(rreg, ltype, line)
                 op_id = ComparisonOp.signed_ops[op]
                 if is_fixed_width_rtype(rreg.type):
@@ -1387,13 +1386,13 @@ def binary_op(self, lreg: Value, rreg: Value, op: str, line: int) -> Value:
                     )
                 if is_tagged(ltype):
                     return self.fixed_width_int_op(rtype, lreg, rreg, op_id, line)
-                if is_bool_rprimitive(ltype) or is_bit_rprimitive(ltype):
+                if is_bool_or_bit_rprimitive(ltype):
                     lreg = self.coerce(lreg, rtype, line)
                     return self.fixed_width_int_op(rtype, lreg, rreg, op_id, line)
             elif op in ComparisonOp.signed_ops:
                 if is_int_rprimitive(ltype):
                     lreg = self.coerce_int_to_fixed_width(lreg, rtype, line)
-                elif is_bool_rprimitive(ltype) or is_bit_rprimitive(ltype):
+                elif is_bool_or_bit_rprimitive(ltype):
                     lreg = self.coerce(lreg, rtype, line)
                 op_id = ComparisonOp.signed_ops[op]
                 if isinstance(lreg, Integer):
@@ -1544,7 +1543,7 @@ def compare_tuples(self, lhs: Value, rhs: Value, op: str, line: int = -1) -> Val
             compare = self.binary_op(lhs_item, rhs_item, op, line)
             # Cast to bool if necessary since most types uses comparison returning a object type
             # See generic_ops.py for more information
-            if not (is_bool_rprimitive(compare.type) or is_bit_rprimitive(compare.type)):
+            if not is_bool_or_bit_rprimitive(compare.type):
                 compare = self.primitive_op(bool_op, [compare], line)
             if i < len(lhs.type.types) - 1:
                 branch = Branch(compare, early_stop, check_blocks[i + 1], Branch.BOOL)
@@ -1563,7 +1562,7 @@ def compare_tuples(self, lhs: Value, rhs: Value, op: str, line: int = -1) -> Val
 
     def translate_instance_contains(self, inst: Value, item: Value, op: str, line: int) -> Value:
         res = self.gen_method_call(inst, "__contains__", [item], None, line)
-        if not is_bool_rprimitive(res.type):
+        if not is_bool_or_bit_rprimitive(res.type):
             res = self.primitive_op(bool_op, [res], line)
         if op == "not in":
             res = self.bool_bitwise_op(res, Integer(1, rtype=bool_rprimitive), "^", line)
@@ -1590,7 +1589,7 @@ def unary_not(self, value: Value, line: int) -> Value:
 
     def unary_op(self, value: Value, expr_op: str, line: int) -> Value:
         typ = value.type
-        if is_bool_rprimitive(typ) or is_bit_rprimitive(typ):
+        if is_bool_or_bit_rprimitive(typ):
             if expr_op == "not":
                 return self.unary_not(value, line)
             if expr_op == "+":
@@ -1748,7 +1747,7 @@ def bool_value(self, value: Value) -> Value:
 
         The result type can be bit_rprimitive or bool_rprimitive.
         """
-        if is_bool_rprimitive(value.type) or is_bit_rprimitive(value.type):
+        if is_bool_or_bit_rprimitive(value.type):
             result = value
         elif is_runtime_subtype(value.type, int_rprimitive):
             zero = Integer(0, short_int_rprimitive)

mypyc/test-data/irbuild-bool.test

@@ -422,3 +422,54 @@ L0:
     r1 = extend r0: builtins.bool to builtins.int
     x = r1
     return x
+
+[case testBitToBoolPromotion]
+def bitand(x: float, y: float, z: float) -> bool:
+    b = (x == y) & (x == z)
+    return b
+def bitor(x: float, y: float, z: float) -> bool:
+    b = (x == y) | (x == z)
+    return b
+def bitxor(x: float, y: float, z: float) -> bool:
+    b = (x == y) ^ (x == z)
+    return b
+def invert(x: float, y: float) -> bool:
+    return not(x == y)
+[out]
+def bitand(x, y, z):
+    x, y, z :: float
+    r0, r1 :: bit
+    r2, b :: bool
+L0:
+    r0 = x == y
+    r1 = x == z
+    r2 = r0 & r1
+    b = r2
+    return b
+def bitor(x, y, z):
+    x, y, z :: float
+    r0, r1 :: bit
+    r2, b :: bool
+L0:
+    r0 = x == y
+    r1 = x == z
+    r2 = r0 | r1
+    b = r2
+    return b
+def bitxor(x, y, z):
+    x, y, z :: float
+    r0, r1 :: bit
+    r2, b :: bool
+L0:
+    r0 = x == y
+    r1 = x == z
+    r2 = r0 ^ r1
+    b = r2
+    return b
+def invert(x, y):
+    x, y :: float
+    r0, r1 :: bit
+L0:
+    r0 = x == y
+    r1 = r0 ^ 1
+    return r1