From 0534e6edf1ddca58d51a7bb8429466f946bb5754 Mon Sep 17 00:00:00 2001
From: Francisco <franciscomacallister@gmail.com>
Date: Fri, 2 Oct 2026 11:32:51 -0500
Subject: [PATCH] Skip the len() pass over the last bitmap in
 RoaringTreemap::select

RoaringTreemap::select calls len() on each bitmap, which walks all of its
containers, and then select() on the bitmap that holds the target, which
walks them again. The last bitmap does not need the len() pass, since
RoaringBitmap::select already returns None when n is past its end.

A treemap whose values all fit in a u32 has a single bitmap, so every
select on it walked the whole bitmap before selecting. On a Ryzen 9 7900X
such selects become about 1.4-3x faster at the middle and the end of the
set and up to ~100x faster near the start. A select that lands in the
last of several bitmaps also skips that bitmap's len() pass. The loop
over the other bitmaps is unchanged.

Add a property test that checks select against the iterator.
---
 roaring/src/treemap/inherent.rs | 41 +++++++++++++++++++++++++++++++--
 1 file changed, 39 insertions(+), 2 deletions(-)

diff --git a/roaring/src/treemap/inherent.rs b/roaring/src/treemap/inherent.rs
index 5eef376f..0cea4808 100644
--- a/roaring/src/treemap/inherent.rs
+++ b/roaring/src/treemap/inherent.rs
@@ -570,7 +570,13 @@ impl RoaringTreemap {
     /// assert_eq!(rb.select(3), None);
     /// ```
     pub fn select(&self, mut n: u64) -> Option<u64> {
-        for (&key, bitmap) in &self.map {
+        let mut bitmaps = self.map.iter();
+        while let Some((&key, bitmap)) = bitmaps.next() {
+            if bitmaps.len() == 0 {
+                // The last bitmap needs no `len()`: its `select` returns `None` past its end.
+                let n = u32::try_from(n).ok()?;
+                return bitmap.select(n).map(|index| util::join(key, index));
+            }
             let len = bitmap.len();
             if len > n {
                 return Some(((key as u64) << 32) | bitmap.select(n as u32).unwrap() as u64);
@@ -602,8 +608,9 @@ impl Clone for RoaringTreemap {
 mod test {
     use crate::{RoaringBitmap, RoaringTreemap};
     use alloc::vec::Vec;
-    use proptest::collection::btree_map;
+    use proptest::collection::{btree_map, vec};
     use proptest::prelude::*;
+    use proptest::sample::Index;
 
     proptest! {
         #[test]
@@ -634,6 +641,36 @@ mod test {
             }
         }
 
+        #[test]
+        fn select_returns_the_nth_value(
+            bitmaps in btree_map(0u32..=16, RoaringBitmap::arbitrary(), 0usize..=16),
+            indices in vec(any::<Index>(), 0..=64),
+        ) {
+            // `from_bitmaps` keeps empty bitmaps, which `select` must step over
+            let treemap = RoaringTreemap::from_bitmaps(bitmaps);
+            let values = treemap.iter().collect::<Vec<_>>();
+
+            // the first and last values of every bitmap, plus random positions
+            let mut positions = Vec::new();
+            let mut offset = 0;
+            for (_, bitmap) in treemap.bitmaps() {
+                let len = bitmap.len();
+                if len > 0 {
+                    positions.extend([offset, offset + len - 1]);
+                }
+                offset += len;
+            }
+            if !values.is_empty() {
+                positions.extend(indices.iter().map(|i| i.index(values.len()) as u64));
+            }
+
+            for n in positions {
+                prop_assert_eq!(treemap.select(n), Some(values[n as usize]));
+            }
+            prop_assert_eq!(treemap.select(values.len() as u64), None);
+            prop_assert_eq!(treemap.select(u64::MAX), None);
+        }
+
         #[test]
         fn retain_preserves_values(mut treemap in RoaringTreemap::arbitrary()) {
             let predicate = |x: u64| !x.is_multiple_of(3);