[mypyc] extend get_expr_length to try constant_fold_expr [1/4] (#19930)

Among other things, `get_expr_length` currently supports StrExpr,
BytesExpr, and string-type Final values.

`constant_fold_expr` already has the code to parse values from all of
these, so we can use it to deduplicate code in `get_expr_length`.

It's able to parse a few more Expressions than our existing
implementation, not that any of these are particularly relevant at the
moment, and allows get_expr_length to automatically benefit from more
relevant future constant-folding enhancements.

Author: BobTheBuidler

mypyc/irbuild/for_helpers.py

@@ -12,7 +12,6 @@
 
 from mypy.nodes import (
     ARG_POS,
-    BytesExpr,
     CallExpr,
     DictionaryComprehension,
     Expression,
@@ -24,10 +23,8 @@
     RefExpr,
     SetExpr,
     StarExpr,
-    StrExpr,
     TupleExpr,
     TypeAlias,
-    Var,
 )
 from mypy.types import LiteralType, TupleType, get_proper_type, get_proper_types
 from mypyc.ir.ops import (
@@ -68,6 +65,7 @@
     short_int_rprimitive,
 )
 from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.constant_fold import constant_fold_expr
 from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME
 from mypyc.irbuild.targets import AssignmentTarget, AssignmentTargetTuple
 from mypyc.primitives.dict_ops import (
@@ -1225,8 +1223,9 @@ def gen_cleanup(self) -> None:
 
 
 def get_expr_length(builder: IRBuilder, expr: Expression) -> int | None:
-    if isinstance(expr, (StrExpr, BytesExpr)):
-        return len(expr.value)
+    folded = constant_fold_expr(builder, expr)
+    if isinstance(folded, (str, bytes)):
+        return len(folded)
     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
@@ -1236,14 +1235,6 @@ def get_expr_length(builder: IRBuilder, expr: Expression) -> int | None:
             return other + sum(stars)  # type: ignore [arg-type]
     elif isinstance(expr, StarExpr):
         return get_expr_length(builder, 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

mypyc/test-data/irbuild-tuple.test

@@ -434,6 +434,54 @@ L4:
     a = r1
     return 1
 
+[case testTupleBuiltFromConstantFolding]
+from typing import Final
+
+c: Final = "c"
+
+def f2(val: str) -> str:
+    return val + "f2"
+
+def test() -> None:
+    # `"ab" + c` should constant fold to "abc"
+    a = tuple(f2(x) for x in "ab" + c)
+
+[out]
+def f2(val):
+    val, r0, r1 :: str
+L0:
+    r0 = 'f2'
+    r1 = PyUnicode_Concat(val, r0)
+    return r1
+def test():
+    r0 :: str
+    r1 :: tuple
+    r2 :: native_int
+    r3 :: bit
+    r4, x, r5 :: str
+    r6 :: native_int
+    a :: tuple
+L0:
+    r0 = 'abc'
+    r1 = PyTuple_New(3)
+    r2 = 0
+    goto L2
+L1:
+    r3 = r2 < 3 :: signed
+    if r3 goto L2 else goto L4 :: bool
+L2:
+    r4 = CPyStr_GetItemUnsafe(r0, r2)
+    x = r4
+    r5 = f2(x)
+    CPySequenceTuple_SetItemUnsafe(r1, r2, r5)
+L3:
+    r6 = r2 + 1
+    r2 = r6
+    goto L1
+L4:
+    a = r1
+    return 1
+
 [case testTupleBuiltFromFinalStr]
 from typing import Final