[mypyc] feat: implement try_getting_literal in preallocate helper [1/1]

mypyc/irbuild/for_helpers.py

@@ -9,6 +9,7 @@
 
 from typing import Callable, ClassVar
 
+from mypy.checkexpr import try_getting_literal
 from mypy.nodes import (
  ARG_POS,
  BytesExpr,
@@ -28,7 +29,7 @@
  TypeAlias,
  Var,
 )
-from mypy.types import LiteralType, TupleType, get_proper_type, get_proper_types
+from mypy.types import LiteralType, TupleType, get_proper_type
 from mypyc.ir.ops import (
  ERR_NEVER,
  BasicBlock,
@@ -255,7 +256,7 @@ def sequence_from_generator_preallocate_helper(
  if isinstance(typ, LiteralType)
  else TupleGet(sequence, i, line)
  )
- for i, typ in enumerate(get_proper_types(proper_type.items))
+ for i, typ in enumerate(map(try_getting_literal, proper_type.items))
  ]
  items = list(map(builder.add, get_item_ops))
  sequence = builder.new_tuple(items, line)

mypyc/test-data/irbuild-tuple.test

@@ -741,76 +741,138 @@ L4:
  a = r11
  return 1
 
-[case testTupleBuiltFromFinalFixedLengthTuple]
+[case testTupleBuiltFromModLevelFixedLengthTuple]
 from typing import Final
 
-source: Final = (1, 2, 3)
+def varint() -> int:
+ # this helper lets us break constant folding for this test
+ return 2
+
+source_final: Final = (1, varint(), 3)
+source_var = (1, varint(), 3)
 
 def f(val: int) -> bool:
  return val % 2 == 0
 
-def test() -> None:
- a = tuple(f(x) for x in source)
+def test_final_can_fold() -> None:
+ a = tuple(f(x) for x in source_final)
+def test_var_can_not_fold() -> None:
+ a = tuple(f(x) for x in source_var)
 [out]
+def varint():
+L0:
+ return 4
 def f(val):
  val, r0 :: int
  r1 :: bit
 L0:
  r0 = CPyTagged_Remainder(val, 4)
  r1 = int_eq r0, 0
  return r1
-def test():
+def test_final_can_fold():
  r0 :: tuple[int, int, int]
  r1 :: bool
  r2 :: object
  r3 :: native_int
  r4 :: bit
- r5, r6, r7 :: int
- r8, r9, r10 :: object
- r11, r12 :: tuple
- r13 :: native_int
- r14 :: bit
- r15 :: object
- r16, x :: int
- r17 :: bool
- r18 :: object
- r19 :: native_int
+ r5 :: object
+ r6 :: int
+ r7, r8 :: object
+ r9, r10 :: tuple
+ r11 :: native_int
+ r12 :: bit
+ r13 :: object
+ r14, x :: int
+ r15 :: bool
+ r16 :: object
+ r17 :: native_int
  a :: tuple
 L0:
- r0 = __main__.source :: static
+ r0 = __main__.source_final :: static
  if is_error(r0) goto L1 else goto L2
 L1:
- r1 = raise NameError('value for final name "source" was not set')
+ r1 = raise NameError('value for final name "source_final" was not set')
  unreachable
 L2:
  r2 = box(tuple[int, int, int], r0)
  r3 = PyObject_Size(r2)
  r4 = r3 >= 0 :: signed
- r5 = r0[0]
+ r5 = object 1
  r6 = r0[1]
- r7 = r0[2]
- r8 = box(int, r5)
- r9 = box(int, r6)
- r10 = box(int, r7)
- r11 = PyTuple_Pack(3, r8, r9, r10)
- r12 = PyTuple_New(r3)
- r13 = 0
+ r7 = object 3
+ r8 = box(int, r6)
+ r9 = PyTuple_Pack(3, r5, r8, r7)
+ r10 = PyTuple_New(r3)
+ r11 = 0
 L3:
- r14 = r13 < r3 :: signed
- if r14 goto L4 else goto L6 :: bool
+ r12 = r11 < r3 :: signed
+ if r12 goto L4 else goto L6 :: bool
 L4:
- r15 = CPySequenceTuple_GetItemUnsafe(r11, r13)
- r16 = unbox(int, r15)
- x = r16
- r17 = f(x)
- r18 = box(bool, r17)
- CPySequenceTuple_SetItemUnsafe(r12, r13, r18)
+ r13 = CPySequenceTuple_GetItemUnsafe(r9, r11)
+ r14 = unbox(int, r13)
+ x = r14
+ r15 = f(x)
+ r16 = box(bool, r15)
+ CPySequenceTuple_SetItemUnsafe(r10, r11, r16)
 L5:
- r19 = r13 + 1
- r13 = r19
+ r17 = r11 + 1
+ r11 = r17
  goto L3
 L6:
- a = r12
+ a = r10
+ return 1
+def test_var_can_not_fold():
+ r0 :: dict
+ r1 :: str
+ r2 :: object
+ r3 :: tuple[int, int, int]
+ r4 :: object
+ r5 :: native_int
+ r6 :: bit
+ r7, r8, r9 :: int
+ r10, r11, r12 :: object
+ r13, r14 :: tuple
+ r15 :: native_int
+ r16 :: bit
+ r17 :: object
+ r18, x :: int
+ r19 :: bool
+ r20 :: object
+ r21 :: native_int
+ a :: tuple
+L0:
+ r0 = __main__.globals :: static
+ r1 = 'source_var'
+ r2 = CPyDict_GetItem(r0, r1)
+ r3 = unbox(tuple[int, int, int], r2)
+ r4 = box(tuple[int, int, int], r3)
+ r5 = PyObject_Size(r4)
+ r6 = r5 >= 0 :: signed
+ r7 = r3[0]
+ r8 = r3[1]
+ r9 = r3[2]
+ r10 = box(int, r7)
+ r11 = box(int, r8)
+ r12 = box(int, r9)
+ r13 = PyTuple_Pack(3, r10, r11, r12)
+ r14 = PyTuple_New(r5)
+ r15 = 0
+L1:
+ r16 = r15 < r5 :: signed
+ if r16 goto L2 else goto L4 :: bool
+L2:
+ r17 = CPySequenceTuple_GetItemUnsafe(r13, r15)
+ r18 = unbox(int, r17)
+ x = r18
+ r19 = f(x)
+ r20 = box(bool, r19)
+ CPySequenceTuple_SetItemUnsafe(r14, r15, r20)
+L3:
+ r21 = r15 + 1
+ r15 = r21
+ goto L1
+L4:
+ a = r14
  return 1
 
 [case testTupleBuiltFromVariableLengthTuple]