perf: map Array.mapFinIdx/mapIdx in place#14224
Open
kim-em wants to merge 1 commit into
Open
Conversation
`Array.mapFinIdxM` (and hence `mapFinIdx`, `mapIdxM`, `mapIdx`) ran only its push-based reference implementation: it reads each element with `as[i]`, a borrow that leaves the element shared with the still-live input array. Any in-place work the mapping function attempts on an element therefore copies it, so a column-style operation that touches one entry of each row of an `Array (Array _)` is O(rows * cols) instead of O(rows), and the gap widens without bound as the rows get wider. Give `mapFinIdxM` an `@[implemented_by]` unsafe implementation mirroring `mapMUnsafe`: it nulls each slot (`uset i default`) before invoking `f`, so on a uniquely-owned input each element arrives exclusive (refcount 1) and may be mutated in place. The reference body is unchanged, so every lemma about `mapFinIdxM` still holds; the `@[inline]` is dropped (kept `@[expose]`) because inlining the reference body at call sites would bypass `@[implemented_by]`, just as `mapM` is not `@[inline]`. tests/compile/arrayMapIdxInPlace.lean is a deterministic regression test using `dbgTraceIfShared`: before this change the `mapFinIdx`/`mapIdx` cases each printed one "row shared" trace per row; now they print none, matching `Array.map`. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
github-actions
Bot
added
the
toolchain-available
|
Mathlib CI status (docs):
|
leanprover-bot
added
the
builds-manual
Collaborator
|
Reference manual CI status:
|
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.
This PR gives
Array.mapFinIdxMan@[implemented_by]unsafe implementation so thatmapFinIdx,mapIdxM, andmapIdxmap in place, likeArray.map.Array.map(viamapMUnsafe) nulls each slot before callingf, so on a uniquely-owned input every element arrives exclusive andfcan mutate it without copying.mapFinIdxMhad no such implementation — only its push-based reference body, which readsas[i]as a borrow, leaving each element shared with the still-live input array.This is a silent performance footgun: reaching for
mapFinIdx/mapIdxinstead ofmapjust to get the index turns in-place work into copying. For anArray (Array _), a column-style operation that bumps one entry of every row becomes O(rows · cols) instead of O(rows) — each row is fully copied to change a single entry — and the gap widens without bound as the rows get wider (measured 6.5× → 17× → 21× as row width goes 32 → 256 → 2048).The fix mirrors
mapMUnsafeexactly, threadingi.toNatand supplying the in-bounds proof vialcProof. The reference body ofmapFinIdxMis unchanged, so its lemmas still hold;@[inline]is dropped (and@[expose]kept) because inlining the reference body would bypass@[implemented_by], exactly asmapMis not@[inline].tests/compile/arrayMapIdxInPlace.leanis a deterministic regression test (no timing) usingdbgTraceIfShared: before this change themapFinIdx/mapIdxcases each printed onerow sharedtrace per row; now they print none, matchingArray.map.Draft: validated against a faithful standalone reproduction on v4.32.0-rc1 (stock prints one trace per row, patched prints none, checksums match); the full-stage build and test suite are left to CI.
🤖 Prepared with Claude Code