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 { - 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::(), 0..=64), + ) { + // `from_bitmaps` keeps empty bitmaps, which `select` must step over + let treemap = RoaringTreemap::from_bitmaps(bitmaps); + let values = treemap.iter().collect::>(); + + // 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);