[mypyc] feat: cache len for iterating over immutable types (#19656)

Currently, if a user uses an immutable type as the sequence input for a
for loop, the length is checked once at each iteration which, while
necessary for some container types such as list and dictionaries, is not
necessary for iterating over immutable types tuple, str, and bytes.

This PR modifies the codebase such that the length is only checked at
the first iteration, and reused from there.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: BobTheBuidler

mypyc/ir/rtypes.py

@@ -628,7 +628,19 @@ def is_range_rprimitive(rtype: RType) -> bool:
 
 def is_sequence_rprimitive(rtype: RType) -> bool:
     return isinstance(rtype, RPrimitive) and (
-        is_list_rprimitive(rtype) or is_tuple_rprimitive(rtype) or is_str_rprimitive(rtype)
+        is_list_rprimitive(rtype)
+        or is_tuple_rprimitive(rtype)
+        or is_str_rprimitive(rtype)
+        or is_bytes_rprimitive(rtype)
+    )
+
+
+def is_immutable_rprimitive(rtype: RType) -> TypeGuard[RPrimitive]:
+    return (
+        is_str_rprimitive(rtype)
+        or is_bytes_rprimitive(rtype)
+        or is_tuple_rprimitive(rtype)
+        or is_frozenset_rprimitive(rtype)
     )
 
 

mypyc/irbuild/builder.py

@@ -91,6 +91,7 @@
     RType,
     RUnion,
     bitmap_rprimitive,
+    bytes_rprimitive,
     c_pyssize_t_rprimitive,
     dict_rprimitive,
     int_rprimitive,
@@ -962,8 +963,12 @@ def get_sequence_type_from_type(self, target_type: Type) -> RType:
         elif isinstance(target_type, Instance):
             if target_type.type.fullname == "builtins.str":
                 return str_rprimitive
-            else:
+            elif target_type.type.fullname == "builtins.bytes":
+                return bytes_rprimitive
+            try:
                 return self.type_to_rtype(target_type.args[0])
+            except IndexError:
+                raise ValueError(f"{target_type!r} is not a valid sequence.") from None
         # This elif-blocks are needed for iterating over classes derived from NamedTuple.
         elif isinstance(target_type, TypeVarLikeType):
             return self.get_sequence_type_from_type(target_type.upper_bound)

mypyc/irbuild/for_helpers.py

@@ -48,6 +48,7 @@
     int_rprimitive,
     is_dict_rprimitive,
     is_fixed_width_rtype,
+    is_immutable_rprimitive,
     is_list_rprimitive,
     is_sequence_rprimitive,
     is_short_int_rprimitive,
@@ -205,9 +206,9 @@ def sequence_from_generator_preallocate_helper(
     there is no condition list in the generator and only one original sequence with
     one index is allowed.
 
-    e.g.  (1) tuple(f(x) for x in a_list/a_tuple)
-          (2) list(f(x) for x in a_list/a_tuple)
-          (3) [f(x) for x in a_list/a_tuple]
+    e.g.  (1) tuple(f(x) for x in a_list/a_tuple/a_str/a_bytes)
+          (2) list(f(x) for x in a_list/a_tuple/a_str/a_bytes)
+          (3) [f(x) for x in a_list/a_tuple/a_str/a_bytes]
     RTuple as an original sequence is not supported yet.
 
     Args:
@@ -224,7 +225,7 @@ def sequence_from_generator_preallocate_helper(
     """
     if len(gen.sequences) == 1 and len(gen.indices) == 1 and len(gen.condlists[0]) == 0:
         rtype = builder.node_type(gen.sequences[0])
-        if is_list_rprimitive(rtype) or is_tuple_rprimitive(rtype) or is_str_rprimitive(rtype):
+        if is_sequence_rprimitive(rtype):
             sequence = builder.accept(gen.sequences[0])
             length = builder.builder.builtin_len(sequence, gen.line, use_pyssize_t=True)
             target_op = empty_op_llbuilder(length, gen.line)
@@ -785,17 +786,31 @@ class ForSequence(ForGenerator):
     Supports iterating in both forward and reverse.
     """
 
+    length_reg: Value | AssignmentTarget | None
+
     def init(self, expr_reg: Value, target_type: RType, reverse: bool) -> None:
+        assert is_sequence_rprimitive(expr_reg.type), expr_reg
         builder = self.builder
         self.reverse = reverse
         # Define target to contain the expression, along with the index that will be used
         # for the for-loop. If we are inside of a generator function, spill these into the
         # environment class.
         self.expr_target = builder.maybe_spill(expr_reg)
+        if is_immutable_rprimitive(expr_reg.type):
+            # If the expression is an immutable type, we can load the length just once.
+            self.length_reg = builder.maybe_spill(self.load_len(self.expr_target))
+        else:
+            # Otherwise, even if the length is known, we must recalculate the length
+            # at every iteration for compatibility with python semantics.
+            self.length_reg = None
         if not reverse:
             index_reg: Value = Integer(0, c_pyssize_t_rprimitive)
         else:
-            index_reg = builder.builder.int_sub(self.load_len(self.expr_target), 1)
+            if self.length_reg is not None:
+                len_val = builder.read(self.length_reg)
+            else:
+                len_val = self.load_len(self.expr_target)
+            index_reg = builder.builder.int_sub(len_val, 1)
         self.index_target = builder.maybe_spill_assignable(index_reg)
         self.target_type = target_type
 
@@ -814,9 +829,13 @@ def gen_condition(self) -> None:
             second_check = BasicBlock()
             builder.add_bool_branch(comparison, second_check, self.loop_exit)
             builder.activate_block(second_check)
-        # For compatibility with python semantics we recalculate the length
-        # at every iteration.
-        len_reg = self.load_len(self.expr_target)
+        if self.length_reg is None:
+            # For compatibility with python semantics we recalculate the length
+            # at every iteration.
+            len_reg = self.load_len(self.expr_target)
+        else:
+            # (unless input is immutable type).
+            len_reg = builder.read(self.length_reg, line)
         comparison = builder.binary_op(builder.read(self.index_target, line), len_reg, "<", line)
         builder.add_bool_branch(comparison, self.body_block, self.loop_exit)
 

mypyc/irbuild/specialize.py

@@ -288,7 +288,7 @@ def translate_tuple_from_generator_call(
     """Special case for simplest tuple creation from a generator.
 
     For example:
-        tuple(f(x) for x in some_list/some_tuple/some_str)
+        tuple(f(x) for x in some_list/some_tuple/some_str/some_bytes)
     'translate_safe_generator_call()' would take care of other cases
     if this fails.
     """

mypyc/test-data/fixtures/ir.py

@@ -172,6 +172,7 @@ def __getitem__(self, i: int) -> int: ...
     def __getitem__(self, i: slice) -> bytes: ...
     def join(self, x: Iterable[object]) -> bytes: ...
     def decode(self, x: str=..., y: str=...) -> str: ...
+    def __iter__(self) -> Iterator[int]: ...
 
 class bytearray:
     @overload

mypyc/test-data/irbuild-generics.test

@@ -711,18 +711,18 @@ L0:
     r0 = __mypyc_self__.__mypyc_env__
     r1 = var_object_size args
     r2 = PyList_New(r1)
-    r3 = 0
+    r3 = var_object_size args
+    r4 = 0
 L1:
-    r4 = var_object_size args
-    r5 = r3 < r4 :: signed
+    r5 = r4 < r3 :: signed
     if r5 goto L2 else goto L4 :: bool
 L2:
-    r6 = CPySequenceTuple_GetItemUnsafe(args, r3)
+    r6 = CPySequenceTuple_GetItemUnsafe(args, r4)
     x = r6
-    CPyList_SetItemUnsafe(r2, r3, x)
+    CPyList_SetItemUnsafe(r2, r4, x)
 L3:
-    r7 = r3 + 1
-    r3 = r7
+    r7 = r4 + 1
+    r4 = r7
     goto L1
 L4:
     can_listcomp = r2