[mypyc] Specialize s[i] == 'x' to a codepoint int compare#21579
Open
VaggelisD wants to merge 1 commit into
Open
[mypyc] Specialize s[i] == 'x' to a codepoint int compare#21579VaggelisD wants to merge 1 commit into
s[i] == 'x' to a codepoint int compare#21579VaggelisD wants to merge 1 commit into
Conversation
Recognizes the AST shape `IndexExpr(str) == StrLiteral` (and the symmetric `StrLiteral == IndexExpr(str)`, plus the `!=` variants) and lowers it to an int compare of codepoints reusing the existing CPyStr_GetItemUnsafeAsInt primitive. Today the pattern lowers to CPyStr_GetItem + CPyStr_EqualLiteral, which allocates or looks up a 1-character PyUnicode object per iteration and goes through a generic string-equality call. After specialization it becomes an inlined PyUnicode_READ plus an int compare -- about 4x faster on bench_str_compare with a 3-compares-per-iteration workload, and closer to ~9x with the more typical 1-compare-per-iteration shape. No annotations required; benefits any code that compares a string index against a 1-character literal. Multi-character / empty literals fall through to the generic path (which still correctly returns False). Bounds checking is preserved -- the helper raises IndexError for out-of-range indices, same as the unspecialized path. Stack: builds on the `ord(s[i])` primitive (python#20578) and the librt.strings codepoint helpers (python#21462, python#21504, python#21509, python#21521, python#21522, python#21553).
p-sawicki
reviewed
Jun 3, 2026
Comment on lines
+1480
to
+1485
| # Going through `Any` routes through the interpreted wrapper, which | ||
| # uses the unspecialized lowering. Confirms the str surface still | ||
| # works for callers that bypass the specializer. | ||
| f: Any = eq_comma | ||
| assert f("hello,world", 5) is True | ||
| assert f("hello", 0) is False |
Collaborator
There was a problem hiding this comment.
i don't think the comment is true, the interpreted wrapper still calls the same generated C function for eq_comma that has the optimization.
to test unspecialized lowering you could add a test case that compares against a one-char str passed as a parameter instead of a literal. i'd imagine we have tests like that already though so i think you could just remove this test case.
p-sawicki
reviewed
Jun 3, 2026
Comment on lines
+994
to
+995
| if isinstance(rhs, IndexExpr) and not isinstance(lhs, IndexExpr): | ||
| lhs, rhs = rhs, lhs |
Collaborator
There was a problem hiding this comment.
i think the errors in the run tests are because of a mypy issue #21586 as it seems rhs is typed as IndexExpr after the swap and assigning lhs to it raises a type error.
you might need to use a temp variable as a work-around as this way it seems to work correctly.
tmp = lhs
lhs, rhs = rhs, tmp
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
7th PR of #21418
Lowers
s[i] == 'x'(and the symmetric==/!=forms) down to a bounds-checked codepoint read + int compare, instead ofCPyStr_GetItem+CPyStr_EqualLiteralwhich (may) allocate a 1-characterPyUnicodeper iteration. No annotations are required for this optimization.On microbenchmarks (1-compare-per-iter hot loop, ~2.5M-codepoint SQL-like string) the comparison is ~3.6x times faster.
Some follow up optimizations that might be worth it I can work on:
s[i] in ('a', 'b', 'c')--> Fuse to one check with N int comparisonss[i] < 'x'--> Need to expand the op set