[mypyc] Use native integers for some sequence indexing operations (#19426)

For example, when iterating over a list, now we use a native integer
for the index (which is not exposed to the user). Previously we used
tagged integers, but in these use cases they provide no real benefit.

This simplifies the IR and should slightly improve performance, as fewer
tagged int to native int conversions are needed.

Multiple ops have to be migrated in one go, as these interact with
each other, and by only changing a subset of them would actually
generate more verbose IR, as a bunch of extra coercions would be
needed.

List of impacted statements:
 * For loop over sequence
 * Assignment like `x, y = a` for tuple/list rvalue
 * Dict iteration
 * List comprehension

For example, consider this example:
```
def foo(a: list[int]) -> None:
    for x in a:
        pass
```

Old generated IR was like this:
```
def foo(a):
    a :: list
    r0 :: short_int
    r1 :: ptr
    r2 :: native_int
    r3 :: short_int
    r4 :: bit
    r5 :: native_int
    r6, r7 :: ptr
    r8 :: native_int
    r9 :: ptr
    r10 :: object
    r11 :: int
    r12 :: short_int
    r13 :: None
L0:
    r0 = 0
L1:
    r1 = get_element_ptr a ob_size :: PyVarObject
    r2 = load_mem r1 :: native_int*
    r3 = r2 << 1
    r4 = r0 < r3 :: signed
    if r4 goto L2 else goto L5 :: bool
L2:
    r5 = r0 >> 1
    r6 = get_element_ptr a ob_item :: PyListObject
    r7 = load_mem r6 :: ptr*
    r8 = r5 * 8
    r9 = r7 + r8
    r10 = load_mem r9 :: builtins.object*
    inc_ref r10
    r11 = unbox(int, r10)
    dec_ref r10
    if is_error(r11) goto L6 (error at foo:2) else goto L3
L3:
    dec_ref r11 :: int
L4:
    r12 = r0 + 2
    r0 = r12
    goto L1
L5:
    return 1
L6:
    r13 = <error> :: None
    return r13
```

Now the generated IR is simpler:
```
def foo(a):
    a :: list
    r0 :: native_int
    r1 :: ptr
    r2 :: native_int
    r3 :: bit
    r4, r5 :: ptr
    r6 :: native_int
    r7 :: ptr
    r8 :: object
    r9 :: int
    r10 :: native_int
    r11 :: None
L0:
    r0 = 0
L1:
    r1 = get_element_ptr a ob_size :: PyVarObject
    r2 = load_mem r1 :: native_int*
    r3 = r0 < r2 :: signed
    if r3 goto L2 else goto L5 :: bool
L2:
    r4 = get_element_ptr a ob_item :: PyListObject
    r5 = load_mem r4 :: ptr*
    r6 = r0 * 8
    r7 = r5 + r6
    r8 = load_mem r7 :: builtins.object*
    inc_ref r8
    r9 = unbox(int, r8)
    dec_ref r8
    if is_error(r9) goto L6 (error at foo:2) else goto L3
L3:
    dec_ref r9 :: int
L4:
    r10 = r0 + 1
    r0 = r10
    goto L1
L5:
    return 1
L6:
    r11 = <error> :: None
    return r11
```

Author: JukkaL

mypyc/irbuild/builder.py

@@ -129,6 +129,7 @@
 from mypyc.primitives.list_ops import list_get_item_unsafe_op, list_pop_last, to_list
 from mypyc.primitives.misc_ops import check_unpack_count_op, get_module_dict_op, import_op
 from mypyc.primitives.registry import CFunctionDescription, function_ops
+from mypyc.primitives.tuple_ops import tuple_get_item_unsafe_op
 
 # These int binary operations can borrow their operands safely, since the
 # primitives take this into consideration.
@@ -772,10 +773,15 @@ def process_sequence_assignment(
         values = []
         for i in range(len(target.items)):
             item = target.items[i]
-            index = self.builder.load_int(i)
+            index: Value
             if is_list_rprimitive(rvalue.type):
+                index = Integer(i, c_pyssize_t_rprimitive)
                 item_value = self.primitive_op(list_get_item_unsafe_op, [rvalue, index], line)
+            elif is_tuple_rprimitive(rvalue.type):
+                index = Integer(i, c_pyssize_t_rprimitive)
+                item_value = self.call_c(tuple_get_item_unsafe_op, [rvalue, index], line)
             else:
+                index = self.builder.load_int(i)
                 item_value = self.builder.gen_method_call(
                     rvalue, "__getitem__", [index], item.type, line
                 )

mypyc/irbuild/for_helpers.py

@@ -44,6 +44,7 @@
     RTuple,
     RType,
     bool_rprimitive,
+    c_pyssize_t_rprimitive,
     int_rprimitive,
     is_dict_rprimitive,
     is_fixed_width_rtype,
@@ -75,6 +76,7 @@
 from mypyc.primitives.misc_ops import stop_async_iteration_op
 from mypyc.primitives.registry import CFunctionDescription
 from mypyc.primitives.set_ops import set_add_op
+from mypyc.primitives.tuple_ops import tuple_get_item_unsafe_op
 
 GenFunc = Callable[[], None]
 
@@ -586,7 +588,9 @@ def gen_cleanup(self) -> None:
 
     def load_len(self, expr: Value | AssignmentTarget) -> Value:
         """A helper to get collection length, used by several subclasses."""
-        return self.builder.builder.builtin_len(self.builder.read(expr, self.line), self.line)
+        return self.builder.builder.builtin_len(
+            self.builder.read(expr, self.line), self.line, use_pyssize_t=True
+        )
 
 
 class ForIterable(ForGenerator):
@@ -766,6 +770,8 @@ def unsafe_index(builder: IRBuilder, target: Value, index: Value, line: int) ->
     # so we just check manually.
     if is_list_rprimitive(target.type):
         return builder.primitive_op(list_get_item_unsafe_op, [target, index], line)
+    elif is_tuple_rprimitive(target.type):
+        return builder.call_c(tuple_get_item_unsafe_op, [target, index], line)
     else:
         return builder.gen_method_call(target, "__getitem__", [index], None, line)
 
@@ -784,11 +790,9 @@ def init(self, expr_reg: Value, target_type: RType, reverse: bool) -> None:
         # environment class.
         self.expr_target = builder.maybe_spill(expr_reg)
         if not reverse:
-            index_reg: Value = Integer(0)
+            index_reg: Value = Integer(0, c_pyssize_t_rprimitive)
         else:
-            index_reg = builder.binary_op(
-                self.load_len(self.expr_target), Integer(1), "-", self.line
-            )
+            index_reg = builder.builder.int_sub(self.load_len(self.expr_target), 1)
         self.index_target = builder.maybe_spill_assignable(index_reg)
         self.target_type = target_type
 
@@ -838,13 +842,7 @@ def gen_step(self) -> None:
         builder = self.builder
         line = self.line
         step = 1 if not self.reverse else -1
-        add = builder.int_op(
-            short_int_rprimitive,
-            builder.read(self.index_target, line),
-            Integer(step),
-            IntOp.ADD,
-            line,
-        )
+        add = builder.builder.int_add(builder.read(self.index_target, line), step)
         builder.assign(self.index_target, add, line)
 
 

mypyc/lib-rt/CPy.h

@@ -652,7 +652,7 @@ PyObject *CPyList_GetItemShortBorrow(PyObject *list, CPyTagged index);
 PyObject *CPyList_GetItemInt64(PyObject *list, int64_t index);
 PyObject *CPyList_GetItemInt64Borrow(PyObject *list, int64_t index);
 bool CPyList_SetItem(PyObject *list, CPyTagged index, PyObject *value);
-bool CPyList_SetItemUnsafe(PyObject *list, CPyTagged index, PyObject *value);
+void CPyList_SetItemUnsafe(PyObject *list, Py_ssize_t index, PyObject *value);
 bool CPyList_SetItemInt64(PyObject *list, int64_t index, PyObject *value);
 PyObject *CPyList_PopLast(PyObject *obj);
 PyObject *CPyList_Pop(PyObject *obj, CPyTagged index);
@@ -703,14 +703,13 @@ tuple_T4CIOO CPyDict_NextItem(PyObject *dict_or_iter, CPyTagged offset);
 int CPyMapping_Check(PyObject *obj);
 
 // Check that dictionary didn't change size during iteration.
-static inline char CPyDict_CheckSize(PyObject *dict, CPyTagged size) {
+static inline char CPyDict_CheckSize(PyObject *dict, Py_ssize_t size) {
     if (!PyDict_CheckExact(dict)) {
         // Dict subclasses will be checked by Python runtime.
         return 1;
     }
-    Py_ssize_t py_size = CPyTagged_AsSsize_t(size);
     Py_ssize_t dict_size = PyDict_Size(dict);
-    if (py_size != dict_size) {
+    if (size != dict_size) {
         PyErr_SetString(PyExc_RuntimeError, "dictionary changed size during iteration");
         return 0;
     }
@@ -783,7 +782,8 @@ bool CPySet_Remove(PyObject *set, PyObject *key);
 
 PyObject *CPySequenceTuple_GetItem(PyObject *tuple, CPyTagged index);
 PyObject *CPySequenceTuple_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end);
-bool CPySequenceTuple_SetItemUnsafe(PyObject *tuple, CPyTagged index, PyObject *value);
+PyObject *CPySequenceTuple_GetItemUnsafe(PyObject *tuple, Py_ssize_t index);
+void CPySequenceTuple_SetItemUnsafe(PyObject *tuple, Py_ssize_t index, PyObject *value);
 
 
 // Exception operations

mypyc/lib-rt/list_ops.c

@@ -231,15 +231,8 @@ bool CPyList_SetItemInt64(PyObject *list, int64_t index, PyObject *value) {
 }
 
 // This function should only be used to fill in brand new lists.
-bool CPyList_SetItemUnsafe(PyObject *list, CPyTagged index, PyObject *value) {
-    if (CPyTagged_CheckShort(index)) {
-        Py_ssize_t n = CPyTagged_ShortAsSsize_t(index);
-        PyList_SET_ITEM(list, n, value);
-        return true;
-    } else {
-        PyErr_SetString(PyExc_OverflowError, CPYTHON_LARGE_INT_ERRMSG);
-        return false;
-    }
+void CPyList_SetItemUnsafe(PyObject *list, Py_ssize_t index, PyObject *value) {
+    PyList_SET_ITEM(list, index, value);
 }
 
 PyObject *CPyList_PopLast(PyObject *obj)

mypyc/lib-rt/tuple_ops.c

@@ -46,16 +46,17 @@ PyObject *CPySequenceTuple_GetSlice(PyObject *obj, CPyTagged start, CPyTagged en
     return CPyObject_GetSlice(obj, start, end);
 }
 
+// No error checking
+PyObject *CPySequenceTuple_GetItemUnsafe(PyObject *tuple, Py_ssize_t index)
+{
+    PyObject *result = PyTuple_GET_ITEM(tuple, index);
+    Py_INCREF(result);
+    return result;
+}
+
 // PyTuple_SET_ITEM does no error checking,
 // and should only be used to fill in brand new tuples.
-bool CPySequenceTuple_SetItemUnsafe(PyObject *tuple, CPyTagged index, PyObject *value)
+void CPySequenceTuple_SetItemUnsafe(PyObject *tuple, Py_ssize_t index, PyObject *value)
 {
-    if (CPyTagged_CheckShort(index)) {
-        Py_ssize_t n = CPyTagged_ShortAsSsize_t(index);
-        PyTuple_SET_ITEM(tuple, n, value);
-        return true;
-    } else {
-        PyErr_SetString(PyExc_OverflowError, CPYTHON_LARGE_INT_ERRMSG);
-        return false;
-    }
+    PyTuple_SET_ITEM(tuple, index, value);
 }

mypyc/primitives/dict_ops.py

@@ -289,7 +289,7 @@
 
 # check that len(dict) == const during iteration
 dict_check_size_op = custom_op(
-    arg_types=[dict_rprimitive, int_rprimitive],
+    arg_types=[dict_rprimitive, c_pyssize_t_rprimitive],
     return_type=bit_rprimitive,
     c_function_name="CPyDict_CheckSize",
     error_kind=ERR_FALSE,

mypyc/primitives/list_ops.py

@@ -13,6 +13,7 @@
     object_rprimitive,
     pointer_rprimitive,
     short_int_rprimitive,
+    void_rtype,
 )
 from mypyc.primitives.registry import (
     ERR_NEG_INT,
@@ -154,7 +155,7 @@
 # that is in-bounds for the list.
 list_get_item_unsafe_op = custom_primitive_op(
     name="list_get_item_unsafe",
-    arg_types=[list_rprimitive, short_int_rprimitive],
+    arg_types=[list_rprimitive, c_pyssize_t_rprimitive],
     return_type=object_rprimitive,
     error_kind=ERR_NEVER,
 )
@@ -183,10 +184,10 @@
 # PyList_SET_ITEM does no error checking,
 # and should only be used to fill in brand new lists.
 new_list_set_item_op = custom_op(
-    arg_types=[list_rprimitive, int_rprimitive, object_rprimitive],
-    return_type=bit_rprimitive,
+    arg_types=[list_rprimitive, c_pyssize_t_rprimitive, object_rprimitive],
+    return_type=void_rtype,
     c_function_name="CPyList_SetItemUnsafe",
-    error_kind=ERR_FALSE,
+    error_kind=ERR_NEVER,
     steals=[False, False, True],
 )
 

mypyc/primitives/tuple_ops.py

@@ -6,14 +6,14 @@
 
 from __future__ import annotations
 
-from mypyc.ir.ops import ERR_FALSE, ERR_MAGIC
+from mypyc.ir.ops import ERR_MAGIC, ERR_NEVER
 from mypyc.ir.rtypes import (
-    bit_rprimitive,
     c_pyssize_t_rprimitive,
     int_rprimitive,
     list_rprimitive,
     object_rprimitive,
     tuple_rprimitive,
+    void_rtype,
 )
 from mypyc.primitives.registry import binary_op, custom_op, function_op, load_address_op, method_op
 
@@ -29,6 +29,15 @@
     error_kind=ERR_MAGIC,
 )
 
+# This is unsafe because it assumes that the index is a non-negative integer
+# that is in-bounds for the tuple.
+tuple_get_item_unsafe_op = custom_op(
+    arg_types=[tuple_rprimitive, c_pyssize_t_rprimitive],
+    return_type=object_rprimitive,
+    c_function_name="CPySequenceTuple_GetItemUnsafe",
+    error_kind=ERR_NEVER,
+)
+
 # Construct a boxed tuple from items: (item1, item2, ...)
 new_tuple_op = custom_op(
     arg_types=[c_pyssize_t_rprimitive],
@@ -48,10 +57,10 @@
 # PyTuple_SET_ITEM does no error checking,
 # and should only be used to fill in brand new tuples.
 new_tuple_set_item_op = custom_op(
-    arg_types=[tuple_rprimitive, int_rprimitive, object_rprimitive],
-    return_type=bit_rprimitive,
+    arg_types=[tuple_rprimitive, c_pyssize_t_rprimitive, object_rprimitive],
+    return_type=void_rtype,
     c_function_name="CPySequenceTuple_SetItemUnsafe",
-    error_kind=ERR_FALSE,
+    error_kind=ERR_NEVER,
     steals=[False, False, True],
 )