[mypyc] Use __index__ for range() arguments in for loops (#22101)

`for i in range(x)` is compiled to a loop that compares the counter with
`x` directly. When an argument isn't statically an `int` (for example
`Any`, `SupportsIndex` or a class with `__index__`), that comparison is
a generic `<` with the original object, but `range()` converts its
arguments with `__index__`. As a result, compiled loops:

- raise `TypeError: '<' not supported between instances of 'int' and
'X'` for objects that define `__index__`, or `TypeError: int object
expected; got X` when such an object is the start value
- accept floats: with `x: Any = 2.5`, `for i in range(x)` loops over 0,
1 and 2 instead of raising `TypeError: 'float' object cannot be
interpreted as an integer`

This converts such arguments once with `PyNumber_Index` before the loop,
like `range()` does, so results and error messages match Python.
Arguments that are already `int`, `bool` or fixed-width ints compile the
same as before.

Since the loop then compares unboxed ints, it's also faster. On Python
3.12, a loop over `range(n)` with `n: Any` goes from about 25 ns to 3 ns
per iteration (2.7 ns with `n: int`).

Author: rheard

mypyc/irbuild/for_helpers.py

@@ -54,13 +54,15 @@
     bool_rprimitive,
     c_pyssize_t_rprimitive,
     int_rprimitive,
+    is_bool_or_bit_rprimitive,
     is_dict_rprimitive,
     is_fixed_width_rtype,
     is_immutable_rprimitive,
     is_list_rprimitive,
     is_sequence_rprimitive,
     is_short_int_rprimitive,
     is_str_rprimitive,
+    is_tagged,
     is_tuple_rprimitive,
     object_pointer_rprimitive,
     object_rprimitive,
@@ -80,7 +82,7 @@
     dict_value_iter_op,
 )
 from mypyc.primitives.exc_ops import no_err_occurred_op, propagate_if_error_op
-from mypyc.primitives.generic_ops import aiter_op, anext_op, iter_op, next_op
+from mypyc.primitives.generic_ops import aiter_op, anext_op, index_op, iter_op, next_op
 from mypyc.primitives.list_ops import (
     list_append_op,
     list_get_item_int64_op,
@@ -1122,6 +1124,8 @@ class ForRange(ForGenerator):
 
     def init(self, start_reg: Value, end_reg: Value, step: int) -> None:
         builder = self.builder
+        start_reg = self.convert_arg(start_reg)
+        end_reg = self.convert_arg(end_reg)
         self.start_reg = start_reg
         self.end_reg = end_reg
         self.step = step
@@ -1139,6 +1143,25 @@ def init(self, start_reg: Value, end_reg: Value, step: int) -> None:
         self.index_target: Register | AssignmentTarget = builder.get_assignment_target(self.index)
         builder.assign(self.index_target, builder.read(self.index_reg, self.line), self.line)
 
+    def convert_arg(self, value: Value) -> Value:
+        """Convert a range() argument to an int using __index__, like range() does."""
+        if (
+            is_tagged(value.type)
+            or is_fixed_width_rtype(value.type)
+            or is_bool_or_bit_rprimitive(value.type)
+        ):
+            return value
+        if (
+            isinstance(value.type, RInstance)
+            and value.type.class_ir.is_ext_class
+            and value.type.class_ir.has_method("__index__")
+        ):
+            # Directly call the __index__ method on native classes that have it.
+            index = self.builder.gen_method_call(value, "__index__", [], int_rprimitive, self.line)
+        else:
+            index = self.builder.call_c(index_op, [value], self.line)
+        return self.builder.coerce(index, int_rprimitive, self.line)
+
     def gen_condition(self) -> None:
         builder = self.builder
         line = self.line

mypyc/primitives/generic_ops.py

@@ -178,6 +178,14 @@
     priority=0,
 )
 
+# operator.index(obj)
+index_op = custom_op(
+    arg_types=[object_rprimitive],
+    return_type=object_rprimitive,
+    c_function_name="PyNumber_Index",
+    error_kind=ERR_MAGIC,
+)
+
 # obj1[obj2]
 py_get_item_op = method_op(
     name="__getitem__",

mypyc/test-data/irbuild-statements.test

@@ -78,6 +78,73 @@ L3:
 L4:
     return 1
 
+[case testForInRangeAnyArgs]
+from typing import Any
+
+def f(a: Any, b: Any) -> None:
+    for i in range(a, b):
+        pass
+[out]
+def f(a, b):
+    a, b, r0 :: object
+    r1 :: int
+    r2 :: object
+    r3, r4, i :: int
+    r5 :: bit
+    r6 :: int
+L0:
+    r0 = PyNumber_Index(a)
+    r1 = unbox(int, r0)
+    r2 = PyNumber_Index(b)
+    r3 = unbox(int, r2)
+    r4 = r1
+    i = r4
+L1:
+    r5 = int_lt r4, r3
+    if r5 goto L2 else goto L4 :: bool
+L2:
+    i = r4
+L3:
+    r6 = CPyTagged_Add(r4, 2)
+    r4 = r6
+    goto L1
+L4:
+    return 1
+
+[case testForInRangeNativeIndexArg]
+class C:
+    def __index__(self) -> int:
+        return 2
+
+def f(c: C) -> None:
+    for i in range(c):  # type: ignore[arg-type]
+        pass
+[out]
+def C.__index__(self):
+    self :: __main__.C
+L0:
+    return 4
+def f(c):
+    c :: __main__.C
+    r0, r1, i :: int
+    r2 :: bit
+    r3 :: int
+L0:
+    r0 = c.__index__()
+    r1 = 0
+    i = r1
+L1:
+    r2 = int_lt r1, r0
+    if r2 goto L2 else goto L4 :: bool
+L2:
+    i = r1
+L3:
+    r3 = CPyTagged_Add(r1, 2)
+    r1 = r3
+    goto L1
+L4:
+    return 1
+
 [case testBreak]
 def f() -> None:
   n = 0

mypyc/test-data/run-loops.test

@@ -572,6 +572,91 @@ def test_range_object() -> None:
     except ValueError as e:
         assert "range() arg 3 must not be zero" in str(e)
 
+[case testRangeLoopIndexArgs]
+from typing import Any
+
+from testutil import assertRaises
+
+def loop_end(end: Any) -> list[int]:
+    result = []
+    for i in range(end):
+        result.append(i)
+    return result
+
+def loop(start: Any, end: Any) -> list[int]:
+    return [i for i in range(start, end)]
+
+def loop_down(start: Any) -> list[int]:
+    return [i for i in range(start, 0, -1)]
+
+def loop_down_to(start: Any, end: Any) -> list[int]:
+    return [i for i in range(start, end, -1)]
+
+def loop_bool(b: bool) -> list[int]:
+    return [i for i in range(b)]
+
+class NativeIndex:
+    def __init__(self, n: int) -> None:
+        self.n = n
+
+    def __index__(self) -> int:
+        return self.n
+
+class NativeIndexSub(NativeIndex):
+    def __index__(self) -> int:
+        return self.n + 1
+
+def loop_native(start: NativeIndex, end: NativeIndex) -> list[int]:
+    # The test fixtures only accept int arguments for range()
+    return [i for i in range(start, end)]  # type: ignore[arg-type]
+
+def test_index_args() -> None:
+    from index_types import Index
+    assert loop_end(Index(3)) == [0, 1, 2]
+    assert loop(Index(1), Index(3)) == [1, 2]
+    assert loop_down(Index(2)) == [2, 1]
+
+def test_big_index_args() -> None:
+    from index_types import Index
+    # Ints >= 2**62 or < -2**62 don't fit in a short int, so these loops cross
+    # between short and heap-allocated values
+    n = 2**62
+    assert loop(Index(n - 1), Index(n + 2)) == [n - 1, n, n + 1]
+    assert loop(Index(-n - 2), Index(-n + 1)) == [-n - 2, -n - 1, -n]
+    assert loop_down_to(Index(n + 1), Index(n - 2)) == [n + 1, n, n - 1]
+    assert loop_down_to(Index(-n + 1), Index(-n - 2)) == [-n + 1, -n, -n - 1]
+    assert loop_end(Index(-n - 1)) == []
+
+def test_native_index_args() -> None:
+    assert loop_native(NativeIndex(1), NativeIndex(3)) == [1, 2]
+    assert loop_native(NativeIndexSub(1), NativeIndex(4)) == [2, 3]
+    n = 2**62
+    assert loop_native(NativeIndex(n - 1), NativeIndex(n + 1)) == [n - 1, n]
+
+def test_int_args() -> None:
+    assert loop_end(3) == [0, 1, 2]
+    assert loop(True, 3) == [1, 2]
+    n = 2**100
+    assert loop(n, n + 2) == [n, n + 1]
+    assert loop_bool(True) == [0]
+    assert loop_bool(False) == []
+
+def test_invalid_args() -> None:
+    with assertRaises(TypeError, "'float' object cannot be interpreted as an integer"):
+        loop_end(2.5)
+    with assertRaises(TypeError, "'str' object cannot be interpreted as an integer"):
+        loop(0, "3")
+    with assertRaises(TypeError, "'float' object cannot be interpreted as an integer"):
+        loop_down(2.0)
+
+[file index_types.py]
+class Index:
+    def __init__(self, n: int) -> None:
+        self.n = n
+
+    def __index__(self) -> int:
+        return self.n
+
 [case testNamedTupleLoop]
 from collections.abc import Iterable
 from typing import NamedTuple, Any