[mypyc] Specialize s[i] == 'x' to a codepoint int compare#21579
Merged
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
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
Contributor
Author
There was a problem hiding this comment.
Interesting, didn't have time to check it out yesterday but that'd definitely confuse me
VaggelisD
force-pushed
the
str-index-compare-specialize
branch
2 times, most recently
from
e116c95 to
60abe82
Compare
- Use a temp variable for the swap normalization. Tuple-unpack form (lhs, rhs = rhs, lhs) interacted badly with mypy's narrowing in mypyc-compiled mypy, producing a runtime IndexExpr-vs-StrExpr cast failure (mypy#21586). Workaround per @p-sawicki on PR python#21579. - Drop test_any_dispatch_uses_generic_path. The 'Any' dispatch still calls the mypyc-compiled eq_comma, which has the specialization, so this test was not exercising the unspecialized path as claimed. The IR golden pins the specialized lowering, and eq_two_chars / eq_empty cover the fall-through behavior.
Rename testStrIndexEqLiteral -> testStrIndexEqLiteral_64bit so it skips on 32-bit. The golden output captures the int-unbox-to-i64 path emitted by translate_getitem_with_bounds_check, which differs on 32-bit (extra 'extend signed i: builtins.int to i64' op shifts register numbering). testOrdOfStrIndex_64bit uses the same primitives and follows the same convention. The fall-through golden (testStrIndexEqLiteralNoSpecialize) keeps no suffix; its IR uses CPyStr_GetItem directly with no unboxing.
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 sets[i] == s[j]--> Need drop the literal-only guard