Skip to content

Optimize WeightedIndex weight lookup and iteration - #1838

Merged
dhardy merged 5 commits into
rust-random:masterfrom
pnn64:perf/weighted-index-iteration
Sep 17, 2026
Merged

dhardy merged 5 commits into
rust-random:masterfrom
pnn64:perf/weighted-index-iteration

Conversation

@pnn64

@pnn64 pnn64 commented Sep 7, 2026

Copy link
Copy Markdown
Contributor

WeightedIndex::weights() can avoid repeated vector growth by reporting its remaining length. Weight lookup can also use a simpler checked access path.

Solution

Implement size_hint and ExactSizeIterator for the weight iterator, allowing collection to reserve enough capacity up front.

Replace the three-way index comparison in weight() with checked slice access and a fallback for the final weight. Returned values, clone/subtraction order, and out-of-bounds behavior are preserved. The sampling algorithm and RNG calls are unchanged.

Tests, benchmarks, and the changelog are updated.

Benchmarks

On a Xeon E5-2696 v4, using rustc 1.98.0 on Windows in release mode.

Each operation below processes 1,024 weights; distribution construction is excluded.

Measurement Before After Change
Collect u32 weights 2,969.77 ns 1,218.95 ns -59.0%
Collect f64 weights 3,175.76 ns 1,507.99 ns -52.5%
Sum u32 weights 1,679.96 ns 1,009.58 ns -39.9%
Sum f64 weights 1,858.66 ns 1,110.69 ns -40.2%
Extend reused buffer, u32 1,848.52 ns 1,129.23 ns -38.9%
Extend reused buffer, f64 1,840.83 ns 1,121.59 ns -39.1%
Growth reallocations per collection 8 0 Eliminated
Initial allocations per collection 1 1 Unchanged

Timing: Criterion, 100 samples, 0.5-second warm-up, and 2-second measurement target. Each version ran twice in before/after/after/before order; times are averages of the two estimates. Allocation counts were measured separately.

Across 36 cases, 27 showed improvements beyond the measured timing uncertainty. The other nine had overlapping before/after confidence-interval envelopes; none showed a resolved regression.

Testing and linting

  • Default library tests: 100 passed.
  • Allocation-only library tests: 84 passed.
  • Iterator coverage includes exact remaining lengths, cloning, partial consumption, and repeated exhaustion.
  • Lookup coverage includes out-of-bounds indices and usize::MAX.
  • Clippy with warnings denied and formatting checks passed.

🤖 AI disclosure: I used Codex with GPT-6 Astra at xhigh reasoning effort to help identify the optimization and construct the benchmark. I have reviewed and understand the resulting implementation.

Simplify checked weight lookup and implement ExactSizeIterator for
weight iteration, avoiding growth reallocations when collecting weights.

Add iterator coverage and benchmarks for collection, lookup, and reuse.

@dhardy dhardy left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the optimisations!

Code looks good to me with some minor suggestions below.

The CHANGELOG needs to link to this PR.

Comment thread src/distr/weighted/weighted_index.rs
Comment thread CHANGELOG.md Outdated
Comment thread benches/benches/weighted.rs
@dhardy
dhardy merged commit f73ce74 into rust-random:master Sep 17, 2026
17 checks passed
@pnn64
pnn64 deleted the perf/weighted-index-iteration branch September 17, 2026 08:13
kodiakhq Bot pushed a commit to pdylanross/fatigue that referenced this pull request Sep 23, 2026
Bumps rand from 0.10.2 to 0.10.3.

Changelog
Sourced from rand's changelog.

[0.10.3] — 2026-09-20
Fixes

Fix WeightedIndex panic when the sum of float weights is infinite; return Error::Overflow instead (#1808)
Fix spurious Error::NonFinite from Uniform::new_inclusive on large finite float ranges such as 0.0..=f64::MAX (#1821)
Fix possible panic due to sampling a deserialized Uniform<char> (#1831)

Changes

Report exact remaining lengths from WeightedIndex::weights() and reduce overhead when reading weights (#1838)

#1808: rust-random/rand#1808
#1821: rust-random/rand#1821
#1831: rust-random/rand#1831
#1838: rust-random/rand#1838



Commits

9e7d328 Prepare rand 0.10.3 (#1840)
f73ce74 Optimize WeightedIndex weight lookup and iteration (#1838)
ef9e044 Avoid panic from deserialized Uniform\<char> where range == 0 (#1831)
c994eb1 docs: fix angle unit in quick start example (#1839)
33dea4f Test that WeightedIndex rejects INFINITY with Error::Overflow (#1822)
94c9078 Fix Uniform::new_inclusive overflow on large finite float ranges (#1821)
bb1262f Use Xoshiro256PlusPlus in examples/rayon-monte-carlo.rs (#1805)
521fab6 Stop pinning dependencies (#1820)
3f7c433 Stop pinning dependencies
cf4f73e sample_efraimidis_spirakis: error on more than amount non-finite weights (#1814)
Additional commits viewable in compare view




Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting @dependabot rebase.


Dependabot commands and options

You can trigger Dependabot actions by commenting on this PR:

@dependabot rebase will rebase this PR
@dependabot recreate will recreate this PR, overwriting any edits that have been made to it
@dependabot show <dependency name> ignore conditions will show all of the ignore conditions of the specified dependency
@dependabot ignore this major version will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
@dependabot ignore this minor version will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
@dependabot ignore this dependency will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants