Optimize WeightedIndex weight lookup and iteration

CHANGELOG.md

@@ -10,6 +10,9 @@ You may also find the [Upgrade Guide](https://rust-random.github.io/book/update.
 
 ## [Unreleased]
 
+### Changes
+- Report exact remaining lengths from `WeightedIndex::weights()` and reduce overhead when reading weights
+
 ### Fixes
 - Fix `WeightedIndex` panic when the sum of float weights is infinite; return `Error::Overflow` instead ([#1808])
 

benches/benches/weighted.rs

@@ -6,8 +6,9 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
-use criterion::{Criterion, black_box, criterion_group, criterion_main};
-use rand::distr::weighted::WeightedIndex;
+use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main};
+use rand::distr::uniform::SampleUniform;
+use rand::distr::weighted::{Weight, WeightedIndex};
 use rand::prelude::*;
 use rand::seq::index::sample_weighted;
 
@@ -19,6 +20,9 @@ criterion_group!(
 criterion_main!(benches);
 
 pub fn bench(c: &mut Criterion) {
+    bench_weight_iteration::<u32>(c, "u32");
+    bench_weight_iteration::<f64>(c, "f64");
+
     c.bench_function("weighted_index_creation", |b| {
         let mut rng = rand::rng();
         let weights = black_box([1u32, 2, 4, 0, 5, 1, 7, 1, 2, 3, 4, 5, 6, 7]);
@@ -58,3 +62,56 @@ pub fn bench(c: &mut Criterion) {
         });
     }
 }
+
+fn bench_weight_iteration<X>(c: &mut Criterion, name: &str)
+where
+    X: SampleUniform + Weight + PartialOrd + From<u32> + core::iter::Sum + for<'a> core::ops::SubAssign<&'a X>,
+{
+    let mut group = c.benchmark_group(format!("weighted_iter/{name}"));
+    for length in [1usize, 4, 16, 64, 256, 1024, 16384] {
+        let distr = WeightedIndex::new((0..length).map(|i| X::from((1 + i % 10) as u32))).unwrap();
+        group.bench_function(BenchmarkId::new("collect", length), |b| {
+            b.iter(|| black_box(&distr).weights().collect::<Vec<_>>())
+        });
+
+        // Control cases: neither summing nor reusing capacity needs a size hint.
+        if [4, 1024].contains(&length) {
+            group
+                .bench_function(BenchmarkId::new("sum", length), |b| b.iter(|| black_box(&distr).weights().sum::<X>()));
+            let mut buffer = Vec::with_capacity(length);
+            group.bench_function(BenchmarkId::new("reuse", length), |b| {
+                b.iter(|| {
+                    buffer.clear();
+                    buffer.extend(black_box(&distr).weights());
+                    black_box(buffer.as_slice());
+                })
+            });
+        }
+
+        if length == 1024 {
+            for (position, index) in [
+                ("first", 0),
+                ("middle", length / 2),
+                ("last", length - 1),
+                ("past_end", length),
+                ("max_index", usize::MAX),
+            ] {
+                group.bench_function(BenchmarkId::new("weight", position), |b| {
+                    b.iter(|| black_box(&distr).weight(black_box(index)))
+                });
+            }
+            let mut iter = distr.weights();
+            let _ = iter.nth(length / 2 - 1);
+            group.bench_function(BenchmarkId::new("collect_remaining", length / 2), |b| {
+                b.iter(|| black_box(iter.clone()).collect::<Vec<_>>())
+            });
+        }
+    }
+    let distr = WeightedIndex::new([X::from(1)]).unwrap();
+    let mut exhausted = distr.weights();
+    let _ = exhausted.next();
+    group.bench_function(BenchmarkId::new("collect", 0), |b| {
+        b.iter(|| black_box(exhausted.clone()).collect::<Vec<_>>())
+    });
+    group.finish();
+}

src/distr/weighted/weighted_index.rs

@@ -289,6 +289,16 @@ where
             }
         }
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let remaining = self.weighted_index.cumulative_weights.len() + 1 - self.index;
+        (remaining, Some(remaining))
+    }
+}
+
+impl<X> ExactSizeIterator for WeightedIndexIter<'_, X> where
+    X: for<'b> core::ops::SubAssign<&'b X> + SampleUniform + PartialOrd + Clone
+{
 }
 
 impl<X: SampleUniform + PartialOrd + Clone> WeightedIndex<X> {
@@ -312,12 +322,12 @@ impl<X: SampleUniform + PartialOrd + Clone> WeightedIndex<X> {
     where
         X: for<'a> core::ops::SubAssign<&'a X>,
     {
-        use core::cmp::Ordering::*;
-
-        let mut weight = match index.cmp(&self.cumulative_weights.len()) {
-            Less => self.cumulative_weights[index].clone(),
-            Equal => self.total_weight.clone(),
-            Greater => return None,
+        let mut weight = if let Some(weight) = self.cumulative_weights.get(index) {
+            weight.clone()
+        } else if index == self.cumulative_weights.len() {
+            self.total_weight.clone()
+        } else {
+            return None;
         };
 
         if index > 0 {
@@ -568,6 +578,7 @@ mod test {
                 assert_eq!(distr.weight(i), Some(*weight));
             }
             assert_eq!(distr.weight(weights.len()), None);
+            assert_eq!(distr.weight(usize::MAX), None);
         }
     }
 
@@ -583,6 +594,20 @@ mod test {
         for weights in data.iter() {
             let distr = WeightedIndex::new(weights.to_vec()).unwrap();
             assert_eq!(distr.weights().collect::<Vec<_>>(), weights.to_vec());
+
+            let mut iter = distr.weights();
+            for (index, expected) in weights.iter().enumerate() {
+                let remaining = weights.len() - index;
+                assert_eq!(iter.size_hint(), (remaining, Some(remaining)));
+                assert_eq!(iter.len(), remaining);
+                assert_eq!(iter.clone().collect::<Vec<_>>(), weights[index..]);
+                assert_eq!(iter.next(), Some(*expected));
+            }
+            for _ in 0..2 {
+                assert_eq!(iter.size_hint(), (0, Some(0)));
+                assert_eq!(iter.len(), 0);
+                assert_eq!(iter.next(), None);
+            }
         }
     }