feat: len only once

Author: BobTheBuidler

mypyc/ir/rtypes.py

@@ -635,6 +635,15 @@ def is_sequence_rprimitive(rtype: RType) -> TypeGuard[RPrimitive]:
     )
 
 
+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)
+    )
+
+
 class TupleNameVisitor(RTypeVisitor[str]):
     """Produce a tuple name based on the concrete representations of types."""
 

mypyc/irbuild/for_helpers.py

@@ -11,15 +11,19 @@
 
 from mypy.nodes import (
     ARG_POS,
+    BytesExpr,
     CallExpr,
     DictionaryComprehension,
     Expression,
     GeneratorExpr,
+    ListExpr,
     Lvalue,
     MemberExpr,
     NameExpr,
+    Optional,
     RefExpr,
     SetExpr,
+    StrExpr,
     TupleExpr,
     TypeAlias,
 )
@@ -46,9 +50,9 @@
     bool_rprimitive,
     c_pyssize_t_rprimitive,
     int_rprimitive,
-    is_bytes_rprimitive,
     is_dict_rprimitive,
     is_fixed_width_rtype,
+    is_immutable_rprimitive,
     is_list_rprimitive,
     is_sequence_rprimitive,
     is_short_int_rprimitive,
@@ -152,6 +156,7 @@ def for_loop_helper_with_index(
     expr_reg: Value,
     body_insts: Callable[[Value], None],
     line: int,
+    length: Value,
 ) -> None:
     """Generate IR for a sequence iteration.
 
@@ -172,7 +177,7 @@ def for_loop_helper_with_index(
     exit_block = BasicBlock()
     condition_block = BasicBlock()
 
-    for_gen = ForSequence(builder, index, body_block, exit_block, line, False)
+    for_gen = ForSequence(builder, index, body_block, exit_block, line, False, length)
     for_gen.init(expr_reg, target_type, reverse=False)
 
     builder.push_loop_stack(step_block, exit_block)
@@ -227,15 +232,15 @@ def sequence_from_generator_preallocate_helper(
         rtype = builder.node_type(gen.sequences[0])
         if is_sequence_rprimitive(rtype):
             sequence = builder.accept(gen.sequences[0])
-            length = builder.builder.builtin_len(sequence, gen.line, use_pyssize_t=True)
+            length = get_expr_length_value(builder, gen.sequences[0], sequence, gen.line, use_pyssize_t=True)
             target_op = empty_op_llbuilder(length, gen.line)
 
             def set_item(item_index: Value) -> None:
                 e = builder.accept(gen.left_expr)
                 builder.call_c(set_item_op, [target_op, item_index, e], gen.line)
 
             for_loop_helper_with_index(
-                builder, gen.indices[0], gen.sequences[0], sequence, set_item, gen.line
+                builder, gen.indices[0], gen.sequences[0], sequence, set_item, gen.line, length
             )
 
             return target_op
@@ -788,20 +793,34 @@ class ForSequence(ForGenerator):
 
     length_reg: Value | AssignmentTarget | None
 
+    def __init__(
+        self,
+        builder: IRBuilder,
+        index: Lvalue,
+        body_block: BasicBlock,
+        loop_exit: BasicBlock,
+        line: int,
+        nested: bool,
+        length: Value | None = None,
+    ) -> None:
+        super().__init__(builder, index, body_block, loop_exit, line, nested)
+        self.length = length
+        """A Value representing the length of the sequence, if known."""
+
     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_tuple_rprimitive(expr_reg.type)
-            or is_str_rprimitive(expr_reg.type)
-            or is_bytes_rprimitive(expr_reg.type)
-        ):
-            self.length_reg = builder.maybe_spill(self.load_len(self.expr_target))
+        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.length or 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)
@@ -1166,3 +1185,24 @@ def gen_step(self) -> None:
     def gen_cleanup(self) -> None:
         for gen in self.gens:
             gen.gen_cleanup()
+
+
+def get_expr_length(expr: Expression) -> Optional[int]:
+    if isinstance(expr, (StrExpr, BytesExpr)):
+        return len(expr.value)
+    elif isinstance(expr, (ListExpr, TupleExpr)):
+        if all(get_expr_length(i) is not None for i in expr.items):
+            return len(expr.items)
+    # TODO: extend this, unrolling should come with a good performance boost
+    return None
+
+
+def get_expr_length_value(builder: IRBuilder, expr: Expression, expr_reg, line: int, use_pyssize_t: bool) -> Value:
+    rtype = builder.node_type(expr)
+    assert is_sequence_rprimitive(rtype), rtype
+    length = get_expr_length(expr)
+    if length is None:
+        # We cannot compute the length at compile time, so we will fetch it.
+        return builder.builder.builtin_len(expr_reg, line, use_pyssize_t=use_pyssize_t)
+    # The expression result is known at compile time, so we can use a constant.
+    return Integer(length, c_pyssize_t_rprimitive if use_pyssize_t else short_int_rprimitive)

mypyc/test-data/fixtures/ir.py

@@ -171,6 +171,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

@@ -689,94 +689,93 @@ def inner_deco_obj.__call__(__mypyc_self__, args, kwargs):
     r0 :: __main__.deco_env
     r1 :: native_int
     r2 :: list
-    r3, r4 :: native_int
-    r5 :: bit
-    r6, x :: object
-    r7 :: native_int
+    r3 :: native_int
+    r4 :: bit
+    r5, x :: object
+    r6 :: native_int
     can_listcomp :: list
-    r8 :: dict
-    r9 :: short_int
-    r10 :: native_int
-    r11 :: object
-    r12 :: tuple[bool, short_int, object, object]
-    r13 :: short_int
-    r14 :: bool
-    r15, r16 :: object
-    r17, k :: str
+    r7 :: dict
+    r8 :: short_int
+    r9 :: native_int
+    r10 :: object
+    r11 :: tuple[bool, short_int, object, object]
+    r12 :: short_int
+    r13 :: bool
+    r14, r15 :: object
+    r16, k :: str
     v :: object
-    r18 :: i32
-    r19, r20, r21 :: bit
+    r17 :: i32
+    r18, r19, r20 :: bit
     can_dictcomp :: dict
-    r22, can_iter, r23, can_use_keys, r24, can_use_values :: list
-    r25 :: object
-    r26 :: list
-    r27 :: object
-    r28 :: dict
-    r29 :: i32
-    r30 :: bit
-    r31 :: tuple
-    r32 :: object
-    r33 :: int
+    r21, can_iter, r22, can_use_keys, r23, can_use_values :: list
+    r24 :: object
+    r25 :: list
+    r26 :: object
+    r27 :: dict
+    r28 :: i32
+    r29 :: bit
+    r30 :: tuple
+    r31 :: object
+    r32 :: int
 L0:
     r0 = __mypyc_self__.__mypyc_env__
     r1 = var_object_size args
     r2 = PyList_New(r1)
-    r3 = var_object_size args
-    r4 = 0
+    r3 = 0
 L1:
-    r5 = r4 < r3 :: signed
-    if r5 goto L2 else goto L4 :: bool
+    r4 = r3 < r1 :: signed
+    if r4 goto L2 else goto L4 :: bool
 L2:
-    r6 = CPySequenceTuple_GetItemUnsafe(args, r4)
-    x = r6
-    CPyList_SetItemUnsafe(r2, r4, x)
+    r5 = CPySequenceTuple_GetItemUnsafe(args, r3)
+    x = r5
+    CPyList_SetItemUnsafe(r2, r3, x)
 L3:
-    r7 = r4 + 1
-    r4 = r7
+    r6 = r3 + 1
+    r3 = r6
     goto L1
 L4:
     can_listcomp = r2
-    r8 = PyDict_New()
-    r9 = 0
-    r10 = PyDict_Size(kwargs)
-    r11 = CPyDict_GetItemsIter(kwargs)
+    r7 = PyDict_New()
+    r8 = 0
+    r9 = PyDict_Size(kwargs)
+    r10 = CPyDict_GetItemsIter(kwargs)
 L5:
-    r12 = CPyDict_NextItem(r11, r9)
-    r13 = r12[1]
-    r9 = r13
-    r14 = r12[0]
-    if r14 goto L6 else goto L8 :: bool
+    r11 = CPyDict_NextItem(r10, r8)
+    r12 = r11[1]
+    r8 = r12
+    r13 = r11[0]
+    if r13 goto L6 else goto L8 :: bool
 L6:
-    r15 = r12[2]
-    r16 = r12[3]
-    r17 = cast(str, r15)
-    k = r17
-    v = r16
-    r18 = CPyDict_SetItem(r8, k, v)
-    r19 = r18 >= 0 :: signed
+    r14 = r11[2]
+    r15 = r11[3]
+    r16 = cast(str, r14)
+    k = r16
+    v = r15
+    r17 = CPyDict_SetItem(r7, k, v)
+    r18 = r17 >= 0 :: signed
 L7:
-    r20 = CPyDict_CheckSize(kwargs, r10)
+    r19 = CPyDict_CheckSize(kwargs, r9)
     goto L5
 L8:
-    r21 = CPy_NoErrOccurred()
+    r20 = CPy_NoErrOccurred()
 L9:
-    can_dictcomp = r8
-    r22 = PySequence_List(kwargs)
-    can_iter = r22
-    r23 = CPyDict_Keys(kwargs)
-    can_use_keys = r23
-    r24 = CPyDict_Values(kwargs)
-    can_use_values = r24
-    r25 = r0.func
-    r26 = PyList_New(0)
-    r27 = CPyList_Extend(r26, args)
-    r28 = PyDict_New()
-    r29 = CPyDict_UpdateInDisplay(r28, kwargs)
-    r30 = r29 >= 0 :: signed
-    r31 = PyList_AsTuple(r26)
-    r32 = PyObject_Call(r25, r31, r28)
-    r33 = unbox(int, r32)
-    return r33
+    can_dictcomp = r7
+    r21 = PySequence_List(kwargs)
+    can_iter = r21
+    r22 = CPyDict_Keys(kwargs)
+    can_use_keys = r22
+    r23 = CPyDict_Values(kwargs)
+    can_use_values = r23
+    r24 = r0.func
+    r25 = PyList_New(0)
+    r26 = CPyList_Extend(r25, args)
+    r27 = PyDict_New()
+    r28 = CPyDict_UpdateInDisplay(r27, kwargs)
+    r29 = r28 >= 0 :: signed
+    r30 = PyList_AsTuple(r25)
+    r31 = PyObject_Call(r24, r30, r27)
+    r32 = unbox(int, r31)
+    return r32
 def deco(func):
     func :: object
     r0 :: __main__.deco_env
@@ -795,3 +794,4 @@ def f(x):
     x :: int
 L0:
     return x
+