[mypyc] feat: cache len for container creation from expressions with length known at compile time (#19503)

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.

Also, in cases where a simple genexp is the input argument for a tuple,
the length is currently checked one additional time before entering the
iteration (this is done to determine how to size the new tuple). In
those cases, we don't even need a length check at the first iteration
step, and can reuse the result of that first `len` call (or compile-time
determined constant) instead.

Lastly, in cases where a tuple is created from a genexp and the length
of the genexp is knowable at compile time, this PR replaces
PyList_AsTuple with the tuple constructor fast-path.

Author: BobTheBuidler

mypyc/irbuild/for_helpers.py

@@ -11,17 +11,22 @@
 
 from mypy.nodes import (
     ARG_POS,
+    BytesExpr,
     CallExpr,
     DictionaryComprehension,
     Expression,
     GeneratorExpr,
+    ListExpr,
     Lvalue,
     MemberExpr,
     NameExpr,
     RefExpr,
     SetExpr,
+    StarExpr,
+    StrExpr,
     TupleExpr,
     TypeAlias,
+    Var,
 )
 from mypyc.ir.ops import (
     ERR_NEVER,
@@ -152,6 +157,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.
 
@@ -173,7 +179,7 @@ def for_loop_helper_with_index(
     condition_block = BasicBlock()
 
     for_gen = ForSequence(builder, index, body_block, exit_block, line, False)
-    for_gen.init(expr_reg, target_type, reverse=False)
+    for_gen.init(expr_reg, target_type, reverse=False, length=length)
 
     builder.push_loop_stack(step_block, exit_block)
 
@@ -227,15 +233,17 @@ 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,17 +796,21 @@ class ForSequence(ForGenerator):
 
     length_reg: Value | AssignmentTarget | None
 
-    def init(self, expr_reg: Value, target_type: RType, reverse: bool) -> None:
+    def init(
+        self, expr_reg: Value, target_type: RType, reverse: bool, length: Value | None = None
+    ) -> None:
         assert is_sequence_rprimitive(expr_reg.type), expr_reg
         builder = self.builder
+        # Record a Value indicating the length of the sequence, if known at compile time.
+        self.length = length
         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))
+            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.
@@ -1166,3 +1178,43 @@ def gen_step(self) -> None:
     def gen_cleanup(self) -> None:
         for gen in self.gens:
             gen.gen_cleanup()
+
+
+def get_expr_length(expr: Expression) -> int | None:
+    if isinstance(expr, (StrExpr, BytesExpr)):
+        return len(expr.value)
+    elif isinstance(expr, (ListExpr, TupleExpr)):
+        # if there are no star expressions, or we know the length of them,
+        # we know the length of the expression
+        stars = [get_expr_length(i) for i in expr.items if isinstance(i, StarExpr)]
+        if None not in stars:
+            other = sum(not isinstance(i, StarExpr) for i in expr.items)
+            return other + sum(stars)  # type: ignore [arg-type]
+    elif isinstance(expr, StarExpr):
+        return get_expr_length(expr.expr)
+    elif (
+        isinstance(expr, RefExpr)
+        and isinstance(expr.node, Var)
+        and expr.node.is_final
+        and isinstance(expr.node.final_value, str)
+        and expr.node.has_explicit_value
+    ):
+        return len(expr.node.final_value)
+    # TODO: extend this, passing length of listcomp and genexp should have worthwhile
+    # performance boost and can be (sometimes) figured out pretty easily. set and dict
+    # comps *can* be done as well but will need special logic to consider the possibility
+    # of key conflicts. Range, enumerate, zip are all simple logic.
+    return None
+
+
+def get_expr_length_value(
+    builder: IRBuilder, expr: Expression, expr_reg: Value, 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/irbuild-generics.test

@@ -678,84 +678,83 @@ 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 :: dict
-    r27 :: object
-    r28 :: int
+    r21, can_iter, r22, can_use_keys, r23, can_use_values :: list
+    r24 :: object
+    r25 :: dict
+    r26 :: object
+    r27 :: 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 = PyDict_SetItem(r8, k, v)
-    r19 = r18 >= 0 :: signed
+    r14 = r11[2]
+    r15 = r11[3]
+    r16 = cast(str, r14)
+    k = r16
+    v = r15
+    r17 = PyDict_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 = PyDict_Copy(kwargs)
-    r27 = PyObject_Call(r25, args, r26)
-    r28 = unbox(int, r27)
-    return r28
+    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 = PyDict_Copy(kwargs)
+    r26 = PyObject_Call(r24, args, r25)
+    r27 = unbox(int, r26)
+    return r27
 def deco(func):
     func :: object
     r0 :: __main__.deco_env