[mypyc] Replace integer floor division by a power of two with a shift

mypyc/irbuild/expression.py

@@ -408,6 +408,8 @@ def transform_op_expr(builder: IRBuilder, expr: OpExpr) -> Value:
     # Special case some int ops to allow borrowing operands.
     if (is_int_rprimitive(builder.node_type(expr.left))
             and is_int_rprimitive(builder.node_type(expr.right))):
+        if expr.op == '//':
+            expr = try_optimize_int_floor_divide(expr)
         if expr.op in int_borrow_friendly_op:
             borrow_left = is_borrow_friendly_expr(builder, expr.right)
             left = builder.accept(expr.left, can_borrow=borrow_left)
@@ -419,6 +421,17 @@ def transform_op_expr(builder: IRBuilder, expr: OpExpr) -> Value:
     )
 
 
+def try_optimize_int_floor_divide(expr: OpExpr) -> OpExpr:
+    """Replace // with a power of two with a right shift, if possible."""
+    if not isinstance(expr.right, IntExpr):
+        return expr
+    divisor = expr.right.value
+    shift = divisor.bit_length() - 1
+    if 0 < shift < 28 and divisor == (1 << shift):
+        return OpExpr('>>', expr.left, IntExpr(shift))
+    return expr
+
+
 def transform_index_expr(builder: IRBuilder, expr: IndexExpr) -> Value:
     index = expr.index
     base_type = builder.node_type(expr.base)

mypyc/test-data/irbuild-int.test

@@ -117,3 +117,41 @@ L5:
     r8 = x
 L6:
     return r8
+
+[case testIntFloorDivideByPowerOfTwo]
+def divby1(x: int) -> int:
+    return x // 1
+def divby2(x: int) -> int:
+    return x // 2
+def divby3(x: int) -> int:
+    return x // 3
+def divby4(x: int) -> int:
+    return x // 4
+def divby8(x: int) -> int:
+    return x // 8
+[out]
+def divby1(x):
+    x, r0 :: int
+L0:
+    r0 = CPyTagged_FloorDivide(x, 2)
+    return r0
+def divby2(x):
+    x, r0 :: int
+L0:
+    r0 = CPyTagged_Rshift(x, 2)
+    return r0
+def divby3(x):
+    x, r0 :: int
+L0:
+    r0 = CPyTagged_FloorDivide(x, 6)
+    return r0
+def divby4(x):
+    x, r0 :: int
+L0:
+    r0 = CPyTagged_Rshift(x, 4)
+    return r0
+def divby8(x):
+    x, r0 :: int
+L0:
+    r0 = CPyTagged_Rshift(x, 6)
+    return r0

mypyc/test-data/run-integers.test

@@ -172,6 +172,8 @@ assert test_isinstance_int_and_not_bool(True) == False
 assert test_isinstance_int_and_not_bool(1) == True
 
 [case testIntOps]
+from typing import Any
+
 def check_and(x: int, y: int) -> None:
     # eval() can be trusted to calculate expected result
     expected = eval('{} & {}'.format(x, y))
@@ -428,6 +430,22 @@ def test_constant_fold() -> None:
     n64 = -(1 << 64) + int()
     assert n64 == -(1 << 64)
 
+def div_by_2(x: int) -> int:
+    return x // 2
+
+def div_by_3(x: int) -> int:
+    return x // 3
+
+def div_by_4(x: int) -> int:
+    return x // 4
+
+def test_floor_divide_by_literal() -> None:
+    for i in range(-100, 100):
+        i_boxed: Any = i
+        assert div_by_2(i) == i_boxed // int('2')
+        assert div_by_3(i) == i_boxed // int('3')
+        assert div_by_4(i) == i_boxed // int('4')
+
 [case testIntMinMax]
 def test_int_min_max() -> None:
     x: int = 200