[mypyc] feat: support constant folding in translate_index_expr [1/1]

mypyc/ir/ops.py

@@ -1045,10 +1045,17 @@ class TupleGet(RegisterOp):
 
     def __init__(self, src: Value, index: int, line: int = -1, *, borrow: bool = False) -> None:
         super().__init__(line)
+        assert isinstance(
+            src.type, RTuple
+        ), "TupleGet only operates on tuples, not {type(src.type).__name__}"
+        src_len = len(src.type.types)
         self.src = src
         self.index = index
-        assert isinstance(src.type, RTuple), "TupleGet only operates on tuples"
-        assert index >= 0
+        if index < 0:
+            self.index += src_len
+        assert (
+            self.index <= src_len - 1
+        ), f"Index out of range.\nsource type: {src.type}\nindex: {index}"
         self.type = src.type.types[index]
         self.is_borrowed = borrow
 

mypyc/irbuild/expression.py

@@ -590,8 +590,12 @@ def transform_index_expr(builder: IRBuilder, expr: IndexExpr) -> Value:
 
     base = builder.accept(expr.base, can_borrow=can_borrow_base)
 
-    if isinstance(base.type, RTuple) and isinstance(index, IntExpr):
-        return builder.add(TupleGet(base, index.value, expr.line))
+    if isinstance(base.type, RTuple):
+        folded_index = constant_fold_expr(builder, index)
+        if isinstance(folded_index, int):
+            length = len(base.type.types)
+            if -length <= folded_index <= length - 1:
+                return builder.add(TupleGet(base, folded_index, expr.line))
 
     if isinstance(index, SliceExpr):
         value = try_gen_slice_op(builder, base, index)