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

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