diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index db986f3fd9a7..20440d4a26f4 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -28,6 +28,7 @@ TypeAlias, Var, ) +from mypy.types import LiteralType, TupleType, get_proper_type, get_proper_types from mypyc.ir.ops import ( ERR_NEVER, BasicBlock, @@ -36,6 +37,7 @@ IntOp, LoadAddress, LoadErrorValue, + LoadLiteral, LoadMem, MethodCall, RaiseStandardError, @@ -170,7 +172,7 @@ def for_loop_helper_with_index( body_insts: a function that generates the body of the loop. It needs a index as parameter. """ - assert is_sequence_rprimitive(expr_reg.type) + assert is_sequence_rprimitive(expr_reg.type), (expr_reg, expr_reg.type) target_type = builder.get_sequence_type(expr) body_block = BasicBlock() @@ -217,10 +219,9 @@ def sequence_from_generator_preallocate_helper( there is no condition list in the generator and only one original sequence with one index is allowed. - e.g. (1) tuple(f(x) for x in a_list/a_tuple/a_str/a_bytes) - (2) list(f(x) for x in a_list/a_tuple/a_str/a_bytes) - (3) [f(x) for x in a_list/a_tuple/a_str/a_bytes] - RTuple as an original sequence is not supported yet. + e.g. (1) tuple(f(x) for x in a_list/a_tuple/a_str/a_bytes/an_rtuple) + (2) list(f(x) for x in a_list/a_tuple/a_str/a_bytes/an_rtuple) + (3) [f(x) for x in a_list/a_tuple/a_str/a_bytes/an_rtuple] Args: empty_op_llbuilder: A function that can generate an empty sequence op when @@ -235,23 +236,41 @@ def sequence_from_generator_preallocate_helper( implementation. """ if len(gen.sequences) == 1 and len(gen.indices) == 1 and len(gen.condlists[0]) == 0: - rtype = builder.node_type(gen.sequences[0]) - if is_sequence_rprimitive(rtype): - sequence = builder.accept(gen.sequences[0]) - 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, length - ) + line = gen.line + sequence_expr = gen.sequences[0] + rtype = builder.node_type(sequence_expr) + if not (is_sequence_rprimitive(rtype) or isinstance(rtype, RTuple)): + return None + sequence = builder.accept(sequence_expr) + length = get_expr_length_value(builder, sequence_expr, sequence, line, use_pyssize_t=True) + if isinstance(rtype, RTuple): + # If input is RTuple, box it to tuple_rprimitive for generic iteration + # TODO: this can be optimized a bit better with an unrolled ForRTuple helper + proper_type = get_proper_type(builder.types[sequence_expr]) + assert isinstance(proper_type, TupleType), proper_type + + get_item_ops = [ + ( + LoadLiteral(typ.value, object_rprimitive) + if isinstance(typ, LiteralType) + else TupleGet(sequence, i, line) + ) + for i, typ in enumerate(get_proper_types(proper_type.items)) + ] + items = list(map(builder.add, get_item_ops)) + sequence = builder.new_tuple(items, line) + + target_op = empty_op_llbuilder(length, 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], line) + + for_loop_helper_with_index( + builder, gen.indices[0], sequence_expr, sequence, set_item, line, length + ) - return target_op + return target_op return None @@ -804,7 +823,7 @@ class ForSequence(ForGenerator): 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 + assert is_sequence_rprimitive(expr_reg.type), (expr_reg, expr_reg.type) builder = self.builder # Record a Value indicating the length of the sequence, if known at compile time. self.length = length @@ -834,7 +853,6 @@ def init( def gen_condition(self) -> None: builder = self.builder line = self.line - # TODO: Don't reload the length each time when iterating an immutable sequence? if self.reverse: # If we are iterating in reverse order, we obviously need # to check that the index is still positive. Somewhat less @@ -1216,7 +1234,7 @@ 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 + assert is_sequence_rprimitive(rtype) or isinstance(rtype, RTuple), rtype length = get_expr_length(expr) if length is None: # We cannot compute the length at compile time, so we will fetch it. diff --git a/mypyc/test-data/irbuild-tuple.test b/mypyc/test-data/irbuild-tuple.test index 7507b6255740..3613c5f0101d 100644 --- a/mypyc/test-data/irbuild-tuple.test +++ b/mypyc/test-data/irbuild-tuple.test @@ -694,36 +694,50 @@ L0: return r1 def test(): r0, source :: tuple[int, int, int] - r1 :: list - r2, r3, r4 :: object - r5, x :: int - r6 :: bool - r7 :: object - r8 :: i32 - r9, r10 :: bit - r11, a :: tuple + r1 :: object + r2 :: native_int + r3 :: bit + r4, r5, r6 :: int + r7, r8, r9 :: object + r10, r11 :: tuple + r12 :: native_int + r13 :: bit + r14 :: object + r15, x :: int + r16 :: bool + r17 :: object + r18 :: native_int + a :: tuple L0: r0 = (2, 4, 6) source = r0 - r1 = PyList_New(0) - r2 = box(tuple[int, int, int], source) - r3 = PyObject_GetIter(r2) + r1 = box(tuple[int, int, int], source) + r2 = PyObject_Size(r1) + r3 = r2 >= 0 :: signed + r4 = source[0] + r5 = source[1] + r6 = source[2] + r7 = box(int, r4) + r8 = box(int, r5) + r9 = box(int, r6) + r10 = PyTuple_Pack(3, r7, r8, r9) + r11 = PyTuple_New(r2) + r12 = 0 L1: - r4 = PyIter_Next(r3) - if is_error(r4) goto L4 else goto L2 + r13 = r12 < r2 :: signed + if r13 goto L2 else goto L4 :: bool L2: - r5 = unbox(int, r4) - x = r5 - r6 = f(x) - r7 = box(bool, r6) - r8 = PyList_Append(r1, r7) - r9 = r8 >= 0 :: signed + r14 = CPySequenceTuple_GetItemUnsafe(r10, r12) + r15 = unbox(int, r14) + x = r15 + r16 = f(x) + r17 = box(bool, r16) + CPySequenceTuple_SetItemUnsafe(r11, r12, r17) L3: + r18 = r12 + 1 + r12 = r18 goto L1 L4: - r10 = CPy_NoErrOccurred() -L5: - r11 = PyList_AsTuple(r1) a = r11 return 1 @@ -746,42 +760,56 @@ L0: r1 = int_eq r0, 0 return r1 def test(): - r0 :: list - r1 :: tuple[int, int, int] - r2 :: bool - r3, r4, r5 :: object - r6, x :: int - r7 :: bool - r8 :: object - r9 :: i32 - r10, r11 :: bit - r12, a :: tuple -L0: - r0 = PyList_New(0) - r1 = __main__.source :: static - if is_error(r1) goto L1 else goto L2 + 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 + a :: tuple +L0: + r0 = __main__.source :: static + if is_error(r0) goto L1 else goto L2 L1: - r2 = raise NameError('value for final name "source" was not set') + r1 = raise NameError('value for final name "source" was not set') unreachable L2: - r3 = box(tuple[int, int, int], r1) - r4 = PyObject_GetIter(r3) + r2 = box(tuple[int, int, int], r0) + r3 = PyObject_Size(r2) + r4 = r3 >= 0 :: signed + r5 = r0[0] + 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 L3: - r5 = PyIter_Next(r4) - if is_error(r5) goto L6 else goto L4 + r14 = r13 < r3 :: signed + if r14 goto L4 else goto L6 :: bool L4: - r6 = unbox(int, r5) - x = r6 - r7 = f(x) - r8 = box(bool, r7) - r9 = PyList_Append(r0, r8) - r10 = r9 >= 0 :: signed + r15 = CPySequenceTuple_GetItemUnsafe(r11, r13) + r16 = unbox(int, r15) + x = r16 + r17 = f(x) + r18 = box(bool, r17) + CPySequenceTuple_SetItemUnsafe(r12, r13, r18) L5: + r19 = r13 + 1 + r13 = r19 goto L3 L6: - r11 = CPy_NoErrOccurred() -L7: - r12 = PyList_AsTuple(r0) a = r12 return 1