feat: add IndexExpr support to constant_fold_expr [1/1]

mypy/constant_fold.py

@@ -11,9 +11,11 @@
     ComplexExpr,
     Expression,
     FloatExpr,
+    IndexExpr,
     IntExpr,
     NameExpr,
     OpExpr,
+    SliceExpr,
     StrExpr,
     UnaryExpr,
     Var,
@@ -73,6 +75,40 @@ def constant_fold_expr(expr: Expression, cur_mod_id: str) -> ConstantValue | Non
         value = constant_fold_expr(expr.expr, cur_mod_id)
         if value is not None:
             return constant_fold_unary_op(expr.op, value)
+    elif isinstance(expr, IndexExpr):
+        base = constant_fold_expr(expr.base, cur_mod_id)
+        if base is not None:
+            index_expr = expr.index
+            if isinstance(index_expr, SliceExpr):
+                if index_expr.begin_index is None:
+                    begin_index = None
+                else:
+                    begin_index = constant_fold_expr(index_expr.begin_index, cur_mod_id)
+                    if begin_index is None:
+                        return None
+                if index_expr.end_index is None:
+                    end_index = None
+                else:
+                    end_index = constant_fold_expr(index_expr.end_index, cur_mod_id)
+                    if end_index is None:
+                        return None
+                if index_expr.stride is None:
+                    stride = None
+                else:
+                    stride = constant_fold_expr(index_expr.stride, cur_mod_id)
+                    if stride is None:
+                        return None
+                try:
+                    return base[begin_index:end_index:stride]  # type: ignore [index, misc]
+                except Exception:
+                    return None
+
+            index = constant_fold_expr(index_expr, cur_mod_id)
+            if index is not None:
+                try:
+                    return base[index]  # type: ignore [index]
+                except Exception:
+                    return None
     return None
 
 

mypyc/irbuild/constant_fold.py

@@ -10,22 +10,25 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Final, Union
+from typing import TYPE_CHECKING, Callable, Final, Union
 
 from mypy.constant_fold import constant_fold_binary_op, constant_fold_unary_op
 from mypy.nodes import (
     BytesExpr,
     ComplexExpr,
     Expression,
     FloatExpr,
+    IndexExpr,
     IntExpr,
     MemberExpr,
     NameExpr,
     OpExpr,
+    SliceExpr,
     StrExpr,
     UnaryExpr,
     Var,
 )
+from mypyc.ir.ops import Value
 from mypyc.irbuild.util import bytes_from_str
 
 if TYPE_CHECKING:
@@ -74,6 +77,40 @@ def constant_fold_expr(builder: IRBuilder, expr: Expression) -> ConstantValue |
         value = constant_fold_expr(builder, expr.expr)
         if value is not None and not isinstance(value, bytes):
             return constant_fold_unary_op(expr.op, value)
+    elif isinstance(expr, IndexExpr):
+        base = constant_fold_expr(builder, expr.base)
+        if base is not None:
+            index_expr = expr.index
+            if isinstance(index_expr, SliceExpr):
+                if index_expr.begin_index is None:
+                    begin_index = None
+                else:
+                    begin_index = constant_fold_expr(builder, index_expr.begin_index)
+                    if begin_index is None:
+                        return None
+                if index_expr.end_index is None:
+                    end_index = None
+                else:
+                    end_index = constant_fold_expr(builder, index_expr.end_index)
+                    if end_index is None:
+                        return None
+                if index_expr.stride is None:
+                    stride = None
+                else:
+                    stride = constant_fold_expr(builder, index_expr.stride)
+                    if stride is None:
+                        return None
+                try:
+                    return base[begin_index:end_index:stride]  # type: ignore [index, misc]
+                except Exception:
+                    return None
+
+            index = constant_fold_expr(builder, index_expr)
+            if index is not None:
+                try:
+                    return base[index]  # type: ignore [index]
+                except Exception:
+                    return None
     return None
 
 
@@ -95,3 +132,29 @@ def constant_fold_binary_op_extended(
         return left * right
 
     return None
+
+
+def try_constant_fold(builder: IRBuilder, expr: Expression) -> Value | None:
+    """Return the constant value of an expression if possible.
+
+    Return None otherwise.
+    """
+    value = constant_fold_expr(builder, expr)
+    if value is not None:
+        return builder.load_literal_value(value)
+    return None
+
+
+def folding_candidate(
+    transform: Callable[[IRBuilder, Expression], Value | None],
+) -> Callable[[IRBuilder, Expression], Value | None]:
+    """Mark a transform function as a candidate for constant folding.
+
+    Candidate functions will attempt to short-circuit the transformation
+    by constant folding the expression and will only proceed to transform
+    the expression if folding is not possible.
+    """
+    def constant_fold_wrap(builder: IRBuilder, expr: Expression) -> Value | None:
+        folded = try_constant_fold(builder, expr)
+        return folded if folded is not None else transform(builder, expr)
+    return constant_fold_wrap

mypyc/irbuild/expression.py

@@ -83,7 +83,7 @@
 )
 from mypyc.irbuild.ast_helpers import is_borrow_friendly_expr, process_conditional
 from mypyc.irbuild.builder import IRBuilder, int_borrow_friendly_op
-from mypyc.irbuild.constant_fold import constant_fold_expr
+from mypyc.irbuild.constant_fold import constant_fold_expr, folding_candidate, try_constant_fold
 from mypyc.irbuild.for_helpers import (
     comprehension_helper,
     raise_error_if_contains_unreachable_names,
@@ -527,11 +527,8 @@ def translate_cast_expr(builder: IRBuilder, expr: CastExpr) -> Value:
 # Operators
 
 
+@folding_candidate
 def transform_unary_expr(builder: IRBuilder, expr: UnaryExpr) -> Value:
-    folded = try_constant_fold(builder, expr)
-    if folded:
-        return folded
-
     return builder.unary_op(builder.accept(expr.expr), expr.op, expr.line)
 
 
@@ -582,6 +579,7 @@ def try_optimize_int_floor_divide(builder: IRBuilder, expr: OpExpr) -> OpExpr:
     return expr
 
 
+@folding_candidate
 def transform_index_expr(builder: IRBuilder, expr: IndexExpr) -> Value:
     index = expr.index
     base_type = builder.node_type(expr.base)
@@ -604,17 +602,6 @@ def transform_index_expr(builder: IRBuilder, expr: IndexExpr) -> Value:
     )
 
 
-def try_constant_fold(builder: IRBuilder, expr: Expression) -> Value | None:
-    """Return the constant value of an expression if possible.
-
-    Return None otherwise.
-    """
-    value = constant_fold_expr(builder, expr)
-    if value is not None:
-        return builder.load_literal_value(value)
-    return None
-
-
 def try_gen_slice_op(builder: IRBuilder, base: Value, index: SliceExpr) -> Value | None:
     """Generate specialized slice op for some index expressions.
 

mypyc/test-data/irbuild-constant-fold.test

@@ -478,3 +478,28 @@ L0:
     r3 = (-1.5+2j)
     neg_2 = r3
     return 1
+
+[case testIndexExprConstantFolding]
+from typing import Final
+
+long_string: Final = "long string"
+
+def pos_index() -> None:
+    a = long_string[5]
+def neg_index() -> None:
+    a = long_string[-5]
+def slice_index() -> None:
+    a = long_string[5:]
+[out]
+def pos_index():
+    a :: str
+L0:
+    a = "s"
+def neg_index():
+    a :: str
+L0:
+    a = "t"
+def slice_index():
+    a :: str
+L0:
+    a = "string"