[mypyc] feat: cache len for iterating over immutable types

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.
     """