From d161662e4b60ae9bc47d9122852f67423eabe265 Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Thu, 23 Oct 2025 09:06:06 -0400 Subject: [PATCH 01/15] [mypyc] feat: enhance ForSequence with start, stop, step This will create optimized code paths for: - `for x in seq[1:]` - `for x in seq[1:2]` - `for x in seq[1:-1]` - `for x in seq[:]` # TODO: exempt this, its invalid - `for x in seq[:2]` - `for x in seq[:-1]` - `for x in seq[1::-1]` - `for x in seq[1:2:-1]` - `for x in seq[1:-1:-1]` - `for x in seq[:2:-1]` - `for x in seq[:-1:-1]` There is no reason to actually create a sliced sequence from the original sequence when we can just grab the appropriate items from the original sequence. We are already grabbing the same number of items from the sliced sequence anyway. --- mypyc/irbuild/for_helpers.py | 51 +++++++++++++++++++++++++++++++----- 1 file changed, 44 insertions(+), 7 deletions(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 715f5432cd13..e7186ff5cb09 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -16,6 +16,7 @@ DictionaryComprehension, Expression, GeneratorExpr, + IndexExpr, ListExpr, Lvalue, MemberExpr, @@ -67,6 +68,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 ( @@ -437,11 +439,23 @@ def make_for_loop_generator( rtyp = builder.node_type(expr) if is_sequence_rprimitive(rtyp): # Special case "for x in ". - expr_reg = builder.accept(expr) target_type = builder.get_sequence_type(expr) - for_list = ForSequence(builder, index, body_block, loop_exit, line, nested) - for_list.init(expr_reg, target_type, reverse=False) + + if isinstance(expr, IndexExpr) and all( + s is None or isinstance(constant_fold_expr(builder, s), int) + for s in (expr.start, expr.stop, expr.step) + ): + for_list.init( + builder.accept(expr.expr), + target_type, + reverse=False, + start=expr.start, + stop=expr.stop, + step=expr.step, + ) + else: + for_list.init(builder.accept(expr), target_type, reverse=False) return for_list if is_dict_rprimitive(rtyp): @@ -821,13 +835,33 @@ class ForSequence(ForGenerator): length_reg: Value | AssignmentTarget | None def init( - self, expr_reg: Value, target_type: RType, reverse: bool, length: Value | None = None + self, + expr_reg: Value, + target_type: RType, + reverse: bool, + length: Value | None = None, + *, + start: int | None = None, + stop: int | None = None, + step: int | None = None, ) -> None: 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 self.reverse = reverse + + self.start = 0 if start is None else start + assert self.start >= 0, "implement me!" + + self.stop = -1 if stop is None else stop + assert self.stop == -1, "implement me!" + + self.step = 1 if step is None else step + assert self.step and self.step >= 1: + if reverse: + self.step *= -1 + # Define target to contain the expression, along with the index that will be used # for the for-loop. If we are inside of a generator function, spill these into the # environment class. @@ -835,13 +869,16 @@ def init( if is_immutable_rprimitive(expr_reg.type): # If the expression is an immutable type, we can load the length just once. self.length_reg = builder.maybe_spill(self.length or self.load_len(self.expr_target)) + # TODO: if stop != -1 implement a safety check and then set self.stop_reg + # gen_condition will need to read stop_reg if present else: # Otherwise, even if the length is known, we must recalculate the length # at every iteration for compatibility with python semantics. self.length_reg = None if not reverse: - index_reg: Value = Integer(0, c_pyssize_t_rprimitive) + index_reg: Value = Integer(self.start, c_pyssize_t_rprimitive) else: + # TODO implement start logic if self.length_reg is not None: len_val = builder.read(self.length_reg) else: @@ -854,6 +891,7 @@ def gen_condition(self) -> None: builder = self.builder line = self.line if self.reverse: + # TODO implement start stop step # If we are iterating in reverse order, we obviously need # to check that the index is still positive. Somewhat less # obviously we still need to check against the length, @@ -898,8 +936,7 @@ def gen_step(self) -> None: # Step to the next item. builder = self.builder line = self.line - step = 1 if not self.reverse else -1 - add = builder.builder.int_add(builder.read(self.index_target, line), step) + add = builder.builder.int_add(builder.read(self.index_target, line), self.step) builder.assign(self.index_target, add, line) From fcd4264397a0274ff1d11aecb2edf5faa3d9d705 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 23 Oct 2025 13:07:59 +0000 Subject: [PATCH 02/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mypyc/irbuild/for_helpers.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index e7186ff5cb09..c4ed97ba289c 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -850,18 +850,18 @@ def init( # Record a Value indicating the length of the sequence, if known at compile time. self.length = length self.reverse = reverse - + self.start = 0 if start is None else start assert self.start >= 0, "implement me!" - + self.stop = -1 if stop is None else stop assert self.stop == -1, "implement me!" - + self.step = 1 if step is None else step assert self.step and self.step >= 1: if reverse: self.step *= -1 - + # Define target to contain the expression, along with the index that will be used # for the for-loop. If we are inside of a generator function, spill these into the # environment class. From e4835aeafa20f0721b54e1bcd761fcd67c5041f2 Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Fri, 24 Oct 2025 11:45:55 -0400 Subject: [PATCH 03/15] Update for_helpers.py --- mypyc/irbuild/for_helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index c4ed97ba289c..b24933e75ccd 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -858,7 +858,7 @@ def init( assert self.stop == -1, "implement me!" self.step = 1 if step is None else step - assert self.step and self.step >= 1: + assert self.step and self.step >= 1, "this should be unreachable for step None and step 0" if reverse: self.step *= -1 From 7dd362a93f6b4a26faf483a6fd1e0ca6382d7f74 Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Fri, 24 Oct 2025 11:47:59 -0400 Subject: [PATCH 04/15] Update for_helpers.py --- mypyc/irbuild/for_helpers.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index b24933e75ccd..f7aec1b0ae4b 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -442,17 +442,17 @@ def make_for_loop_generator( target_type = builder.get_sequence_type(expr) for_list = ForSequence(builder, index, body_block, loop_exit, line, nested) - if isinstance(expr, IndexExpr) and all( + if isinstance(expr, IndexExpr) and isinstance(expr.index, SliceExpr) and all( s is None or isinstance(constant_fold_expr(builder, s), int) - for s in (expr.start, expr.stop, expr.step) + for s in (expr.index.start, expr.index.stop, expr.index.step) ): for_list.init( - builder.accept(expr.expr), + builder.accept(expr.base), target_type, reverse=False, - start=expr.start, - stop=expr.stop, - step=expr.step, + start=expr.index.start, + stop=expr.index.stop, + step=expr.index.step, ) else: for_list.init(builder.accept(expr), target_type, reverse=False) From 17a5d6d00ab6c3319e171f60671c9096f453c40b Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Fri, 24 Oct 2025 11:48:29 -0400 Subject: [PATCH 05/15] Update for_helpers.py --- mypyc/irbuild/for_helpers.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index f7aec1b0ae4b..57c3cc6fad1a 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -23,6 +23,7 @@ NameExpr, RefExpr, SetExpr, + SliceExpr, StarExpr, StrExpr, TupleExpr, From 2ac9a737c4738fddac8648d8a199ae79a3370514 Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Fri, 24 Oct 2025 11:49:05 -0400 Subject: [PATCH 06/15] Update for_helpers.py --- mypyc/irbuild/for_helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 57c3cc6fad1a..840368f0bade 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -439,7 +439,7 @@ def make_for_loop_generator( rtyp = builder.node_type(expr) if is_sequence_rprimitive(rtyp): - # Special case "for x in ". + # Special case "for x in ". target_type = builder.get_sequence_type(expr) for_list = ForSequence(builder, index, body_block, loop_exit, line, nested) From 62adcfafb881b13d6008578d223194747f3e3d29 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 24 Oct 2025 15:49:31 +0000 Subject: [PATCH 07/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mypyc/irbuild/for_helpers.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 840368f0bade..0e289edda2a3 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -443,9 +443,13 @@ def make_for_loop_generator( target_type = builder.get_sequence_type(expr) for_list = ForSequence(builder, index, body_block, loop_exit, line, nested) - if isinstance(expr, IndexExpr) and isinstance(expr.index, SliceExpr) and all( - s is None or isinstance(constant_fold_expr(builder, s), int) - for s in (expr.index.start, expr.index.stop, expr.index.step) + if ( + isinstance(expr, IndexExpr) + and isinstance(expr.index, SliceExpr) + and all( + s is None or isinstance(constant_fold_expr(builder, s), int) + for s in (expr.index.start, expr.index.stop, expr.index.step) + ) ): for_list.init( builder.accept(expr.base), From 2e1b99959aa115e36ad97cfe2efb664172641b89 Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Fri, 24 Oct 2025 11:50:06 -0400 Subject: [PATCH 08/15] Update for_helpers.py --- mypyc/irbuild/for_helpers.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 0e289edda2a3..0739d18ccc4c 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -443,21 +443,17 @@ def make_for_loop_generator( target_type = builder.get_sequence_type(expr) for_list = ForSequence(builder, index, body_block, loop_exit, line, nested) - if ( - isinstance(expr, IndexExpr) - and isinstance(expr.index, SliceExpr) - and all( - s is None or isinstance(constant_fold_expr(builder, s), int) - for s in (expr.index.start, expr.index.stop, expr.index.step) - ) + if isinstance(expr, IndexExpr) and isinstance(expr.index, SliceExpr) and all( + s is None or isinstance(constant_fold_expr(builder, s), int) + for s in (expr.index.start, expr.index.stop, expr.index.step) ): for_list.init( builder.accept(expr.base), target_type, reverse=False, - start=expr.index.start, - stop=expr.index.stop, - step=expr.index.step, + start=constant_fold_expr(builder, expr.index.start), + stop=constant_fold_expr(builder, expr.index.stop), + step=constant_fold_expr(builder, expr.index.step), ) else: for_list.init(builder.accept(expr), target_type, reverse=False) From 090a440746888d4130f8a90861a8484e3e28a5e9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 24 Oct 2025 15:53:11 +0000 Subject: [PATCH 09/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mypyc/irbuild/for_helpers.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 0739d18ccc4c..2373ade4aefd 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -443,9 +443,13 @@ def make_for_loop_generator( target_type = builder.get_sequence_type(expr) for_list = ForSequence(builder, index, body_block, loop_exit, line, nested) - if isinstance(expr, IndexExpr) and isinstance(expr.index, SliceExpr) and all( - s is None or isinstance(constant_fold_expr(builder, s), int) - for s in (expr.index.start, expr.index.stop, expr.index.step) + if ( + isinstance(expr, IndexExpr) + and isinstance(expr.index, SliceExpr) + and all( + s is None or isinstance(constant_fold_expr(builder, s), int) + for s in (expr.index.start, expr.index.stop, expr.index.step) + ) ): for_list.init( builder.accept(expr.base), From 0cf4c31b89c7855605fe588994c4ec6258fcbe4c Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Fri, 24 Oct 2025 11:54:45 -0400 Subject: [PATCH 10/15] Update for_helpers.py --- mypyc/irbuild/for_helpers.py | 38 ++++++++++++++++++++---------------- 1 file changed, 21 insertions(+), 17 deletions(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 2373ade4aefd..e77e6d3c4ae3 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -443,24 +443,28 @@ def make_for_loop_generator( target_type = builder.get_sequence_type(expr) for_list = ForSequence(builder, index, body_block, loop_exit, line, nested) - if ( - isinstance(expr, IndexExpr) - and isinstance(expr.index, SliceExpr) - and all( + if isinstance(expr, IndexExpr) and isinstance(expr.index, SliceExpr): + # TODO: maybe we must not apply this optimization to list type specifically + # because the need to check length changes at each iteration? + start = expr.index.start + stop = expr.index.stop + step = expr.index.step + + if all( s is None or isinstance(constant_fold_expr(builder, s), int) - for s in (expr.index.start, expr.index.stop, expr.index.step) - ) - ): - for_list.init( - builder.accept(expr.base), - target_type, - reverse=False, - start=constant_fold_expr(builder, expr.index.start), - stop=constant_fold_expr(builder, expr.index.stop), - step=constant_fold_expr(builder, expr.index.step), - ) - else: - for_list.init(builder.accept(expr), target_type, reverse=False) + for s in (start, stop, step) + ): + for_list.init( + builder.accept(expr.base), + target_type, + reverse=False, + start=constant_fold_expr(builder, start), + stop=constant_fold_expr(builder, stop), + step=constant_fold_expr(builder, step), + ) + return for_list + + for_list.init(builder.accept(expr), target_type, reverse=False) return for_list if is_dict_rprimitive(rtyp): From fbf88f60d496a72cb83e858b3202e86b52c9da6f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 24 Oct 2025 15:58:19 +0000 Subject: [PATCH 11/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mypyc/irbuild/for_helpers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index e77e6d3c4ae3..7bd256b05b43 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -449,7 +449,7 @@ def make_for_loop_generator( start = expr.index.start stop = expr.index.stop step = expr.index.step - + if all( s is None or isinstance(constant_fold_expr(builder, s), int) for s in (start, stop, step) @@ -463,7 +463,7 @@ def make_for_loop_generator( step=constant_fold_expr(builder, step), ) return for_list - + for_list.init(builder.accept(expr), target_type, reverse=False) return for_list From 5d48f2e43d3d183d3a0453429e15f18808093f72 Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Fri, 24 Oct 2025 14:24:15 -0400 Subject: [PATCH 12/15] Update for_helpers.py --- mypyc/irbuild/for_helpers.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 7bd256b05b43..1845863f2e9e 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -446,21 +446,22 @@ def make_for_loop_generator( if isinstance(expr, IndexExpr) and isinstance(expr.index, SliceExpr): # TODO: maybe we must not apply this optimization to list type specifically # because the need to check length changes at each iteration? - start = expr.index.start - stop = expr.index.stop - step = expr.index.step - - if all( - s is None or isinstance(constant_fold_expr(builder, s), int) - for s in (start, stop, step) - ): + + def constant_fold_or_none(expr: Expression | None) -> Any: + return None if expr is None else constant_fold_expr(builder, expr) + + start = constant_fold_or_none(expr.index.start) + stop = constant_fold_or_none(expr.index.stop) + step = constant_fold_or_none(expr.index.step) + + if all(s is None or isinstance(s, int) for s in (start, stop, step)): for_list.init( builder.accept(expr.base), target_type, reverse=False, - start=constant_fold_expr(builder, start), - stop=constant_fold_expr(builder, stop), - step=constant_fold_expr(builder, step), + start=start, + stop=stop, + step=step, ) return for_list From 662b030ce0523188b536b44d13016d62d974d58a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 24 Oct 2025 18:25:38 +0000 Subject: [PATCH 13/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mypyc/irbuild/for_helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 1845863f2e9e..69e229e0fe44 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -446,7 +446,7 @@ def make_for_loop_generator( if isinstance(expr, IndexExpr) and isinstance(expr.index, SliceExpr): # TODO: maybe we must not apply this optimization to list type specifically # because the need to check length changes at each iteration? - + def constant_fold_or_none(expr: Expression | None) -> Any: return None if expr is None else constant_fold_expr(builder, expr) From c1f109685f4d2ec50635a9f6fc964b9de3fcf618 Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Fri, 24 Oct 2025 15:00:35 -0400 Subject: [PATCH 14/15] Update for_helpers.py --- mypyc/irbuild/for_helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 69e229e0fe44..93989826ea57 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -7,7 +7,7 @@ from __future__ import annotations -from typing import Callable, ClassVar +from typing import Any, Callable, ClassVar from mypy.nodes import ( ARG_POS, From ff72e93f577c54263c6990d2036a887418867a3e Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Fri, 24 Oct 2025 15:03:52 -0400 Subject: [PATCH 15/15] Update for_helpers.py --- mypyc/irbuild/for_helpers.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 93989826ea57..0ebdf597fdd7 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -450,9 +450,9 @@ def make_for_loop_generator( def constant_fold_or_none(expr: Expression | None) -> Any: return None if expr is None else constant_fold_expr(builder, expr) - start = constant_fold_or_none(expr.index.start) - stop = constant_fold_or_none(expr.index.stop) - step = constant_fold_or_none(expr.index.step) + start = constant_fold_or_none(expr.index.begin_index) + stop = constant_fold_or_none(expr.index.end_index) + step = constant_fold_or_none(expr.index.stride) if all(s is None or isinstance(s, int) for s in (start, stop, step)): for_list.init(