Skip to content

Commit 282e13e

Browse files
Dr-EmannDr-Emann
committed
Introduce Iter::next_range and next_range_back
1 parent 36afadf commit 282e13e

File tree

9 files changed

+1194
-38
lines changed

9 files changed

+1194
-38
lines changed

‎fuzz/Cargo.lock‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎fuzz/fuzz_targets/arbitrary_ops/mod.rs‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,8 @@ pub enum BitmapIteratorOperation {
103103
AdvanceBackTo(Num),
104104
Nth(Num),
105105
NthBack(Num),
106+
NextRange,
107+
NextRangeBack,
106108
}
107109

108110
impl ReadBitmapOperation {
@@ -466,6 +468,69 @@ impl<'a> CRoaringIterRange<'a> {
466468
}
467469
self.next_back()
468470
}
471+
472+
fn next_range(&mut self) -> Option<RangeInclusive<u32>> {
473+
if self.empty {
474+
return None;
475+
}
476+
self.cursor.reset_at_or_after(self.start);
477+
let range_start = self.cursor.current()?;
478+
if range_start > self.end_inclusive {
479+
self.empty = true;
480+
return None;
481+
}
482+
let mut range_end_inclusive = range_start;
483+
while range_end_inclusive < self.end_inclusive {
484+
if let Some(next) = self.cursor.next() {
485+
if next == range_end_inclusive + 1 {
486+
range_end_inclusive = next;
487+
continue;
488+
}
489+
}
490+
break;
491+
}
492+
493+
if range_end_inclusive == self.end_inclusive {
494+
self.empty = true;
495+
} else {
496+
self.start = range_end_inclusive + 1;
497+
}
498+
499+
Some(range_start..=range_end_inclusive)
500+
}
501+
502+
fn next_range_back(&mut self) -> Option<RangeInclusive<u32>> {
503+
if self.empty {
504+
return None;
505+
}
506+
self.cursor.reset_at_or_after(self.end_inclusive);
507+
if self.cursor.current().is_none_or(|n| n > self.end_inclusive) {
508+
self.cursor.move_prev();
509+
}
510+
let range_end_inclusive = self.cursor.current()?;
511+
if range_end_inclusive < self.start {
512+
self.empty = true;
513+
return None;
514+
}
515+
let mut range_start = range_end_inclusive;
516+
while range_start > self.start {
517+
if let Some(prev) = self.cursor.prev() {
518+
if prev == range_start - 1 {
519+
range_start = prev;
520+
continue;
521+
}
522+
}
523+
break;
524+
}
525+
526+
if range_start == self.start {
527+
self.empty = true;
528+
} else {
529+
self.end_inclusive = range_start - 1;
530+
}
531+
532+
Some(range_start..=range_end_inclusive)
533+
}
469534
}
470535

471536
impl BitmapIteratorOperation {
@@ -491,6 +556,12 @@ impl BitmapIteratorOperation {
491556
BitmapIteratorOperation::NthBack(n) => {
492557
assert_eq!(x.nth_back(n.0), y.nth_back(n.0 as usize));
493558
}
559+
BitmapIteratorOperation::NextRange => {
560+
assert_eq!(x.next_range(), y.next_range());
561+
}
562+
BitmapIteratorOperation::NextRangeBack => {
563+
assert_eq!(x.next_range_back(), y.next_range_back());
564+
}
494565
}
495566
}
496567
}

‎roaring/src/bitmap/container.rs‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -438,13 +438,41 @@ impl DoubleEndedIterator for Iter<'_> {
438438
impl ExactSizeIterator for Iter<'_> {}
439439

440440
impl Iter<'_> {
441+
pub(crate) fn peek(&self) -> Option<u32> {
442+
self.inner.peek().map(|i| util::join(self.key, i))
443+
}
444+
445+
pub(crate) fn peek_back(&self) -> Option<u32> {
446+
self.inner.peek_back().map(|i| util::join(self.key, i))
447+
}
448+
441449
pub(crate) fn advance_to(&mut self, index: u16) {
442450
self.inner.advance_to(index);
443451
}
444452

445453
pub(crate) fn advance_back_to(&mut self, index: u16) {
446454
self.inner.advance_back_to(index);
447455
}
456+
457+
/// Returns the range of consecutive set bits from the current position to the end of the current run
458+
///
459+
/// After this call, the iterator will be positioned at the first item after the returned range.
460+
/// Returns `None` if the iterator is exhausted.
461+
pub(crate) fn next_range(&mut self) -> Option<RangeInclusive<u32>> {
462+
self.inner
463+
.next_range()
464+
.map(|r| util::join(self.key, *r.start())..=util::join(self.key, *r.end()))
465+
}
466+
467+
/// Returns the range of consecutive set bits from the start of the current run to the current back position
468+
///
469+
/// After this call, the back of the iterator will be positioned at the last item before the returned range.
470+
/// Returns `None` if the iterator is exhausted.
471+
pub(crate) fn next_range_back(&mut self) -> Option<RangeInclusive<u32>> {
472+
self.inner
473+
.next_range_back()
474+
.map(|r| util::join(self.key, *r.start())..=util::join(self.key, *r.end()))
475+
}
448476
}
449477

450478
impl fmt::Debug for Container {

‎roaring/src/bitmap/iter.rs‎

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,100 @@ fn advance_to_impl<'a, It>(
9292
}
9393
}
9494

95+
fn next_range_impl<'a, It>(
96+
front_iter: &mut Option<container::Iter<'a>>,
97+
containers: &mut It,
98+
back_iter: &mut Option<container::Iter<'a>>,
99+
) -> Option<core::ops::RangeInclusive<u32>>
100+
where
101+
It: Iterator + Clone,
102+
It: AsRef<[Container]>,
103+
It::Item: IntoIterator<IntoIter = container::Iter<'a>>,
104+
{
105+
let range = loop {
106+
if let Some(r) = and_then_or_clear(front_iter, container::Iter::next_range) {
107+
break r;
108+
}
109+
*front_iter = match containers.next() {
110+
Some(inner) => Some(inner.into_iter()),
111+
None => return and_then_or_clear(back_iter, container::Iter::next_range),
112+
}
113+
};
114+
let (range_start, mut range_end) = (*range.start(), *range.end());
115+
while range_end & 0xFFFF == 0xFFFF {
116+
let Some(after_end) = range_end.checked_add(1) else {
117+
return Some(range_start..=range_end);
118+
};
119+
let (next_key, _) = util::split(after_end);
120+
121+
if containers.as_ref().first().is_some_and(|c| c.key == next_key && c.contains(0)) {
122+
let mut iter = containers.next().unwrap().into_iter();
123+
let next_range = iter.next_range().unwrap();
124+
*front_iter = Some(iter);
125+
debug_assert_eq!(*next_range.start(), after_end);
126+
range_end = *next_range.end();
127+
} else {
128+
if let Some(iter) = back_iter {
129+
if iter.peek() == Some(after_end) {
130+
let next_range = iter.next_range().unwrap();
131+
debug_assert_eq!(*next_range.start(), after_end);
132+
range_end = *next_range.end();
133+
}
134+
}
135+
break;
136+
}
137+
}
138+
139+
Some(range_start..=range_end)
140+
}
141+
142+
fn next_range_back_impl<'a, It>(
143+
front_iter: &mut Option<container::Iter<'a>>,
144+
containers: &mut It,
145+
back_iter: &mut Option<container::Iter<'a>>,
146+
) -> Option<core::ops::RangeInclusive<u32>>
147+
where
148+
It: DoubleEndedIterator,
149+
It: AsRef<[Container]>,
150+
It::Item: IntoIterator<IntoIter = container::Iter<'a>>,
151+
{
152+
let range = loop {
153+
if let Some(r) = and_then_or_clear(back_iter, container::Iter::next_range_back) {
154+
break r;
155+
}
156+
*back_iter = match containers.next_back() {
157+
Some(inner) => Some(inner.into_iter()),
158+
None => return and_then_or_clear(front_iter, container::Iter::next_range_back),
159+
}
160+
};
161+
let (mut range_start, range_end) = (*range.start(), *range.end());
162+
while range_start & 0xFFFF == 0 {
163+
let Some(before_start) = range_start.checked_sub(1) else {
164+
return Some(range_start..=range_end);
165+
};
166+
let (prev_key, _) = util::split(before_start);
167+
168+
if containers.as_ref().last().is_some_and(|c| c.key == prev_key && c.contains(u16::MAX)) {
169+
let mut iter = containers.next_back().unwrap().into_iter();
170+
let next_range = iter.next_range_back().unwrap();
171+
*back_iter = Some(iter);
172+
debug_assert_eq!(*next_range.end(), before_start);
173+
range_start = *next_range.start();
174+
} else {
175+
if let Some(iter) = front_iter {
176+
if iter.key == prev_key && iter.peek_back() == Some(before_start) {
177+
let next_range = iter.next_range_back().unwrap();
178+
debug_assert_eq!(*next_range.end(), before_start);
179+
range_start = *next_range.start();
180+
}
181+
}
182+
break;
183+
}
184+
}
185+
186+
Some(range_start..=range_end)
187+
}
188+
95189
fn advance_back_to_impl<'a, It>(
96190
n: u32,
97191
front_iter: &mut Option<container::Iter<'a>>,
@@ -197,6 +291,46 @@ impl Iter<'_> {
197291
pub fn advance_back_to(&mut self, n: u32) {
198292
advance_back_to_impl(n, &mut self.front, &mut self.containers, &mut self.back);
199293
}
294+
295+
/// Returns the range of consecutive set bits from the current position to the end of the current run
296+
///
297+
/// After this call, the iterator will be positioned at the first item after the returned range.
298+
/// Returns `None` if the iterator is exhausted.
299+
///
300+
/// # Examples
301+
///
302+
/// ```rust
303+
/// use roaring::RoaringBitmap;
304+
///
305+
/// let bm = RoaringBitmap::from([1, 2, 4, 5]);
306+
/// let mut iter = bm.iter();
307+
/// assert_eq!(iter.next_range(), Some(1..=2));
308+
/// assert_eq!(iter.next(), Some(4));
309+
/// assert_eq!(iter.next_range(), Some(5..=5));
310+
/// ```
311+
pub fn next_range(&mut self) -> Option<core::ops::RangeInclusive<u32>> {
312+
next_range_impl(&mut self.front, &mut self.containers, &mut self.back)
313+
}
314+
315+
/// Returns the range of consecutive set bits from the start of the current run to the current back position
316+
///
317+
/// After this call, the back of the iterator will be positioned at the last item before the returned range.
318+
/// Returns `None` if the iterator is exhausted.
319+
///
320+
/// # Examples
321+
///
322+
/// ```rust
323+
/// use roaring::RoaringBitmap;
324+
///
325+
/// let bm = RoaringBitmap::from([1, 2, 4, 5]);
326+
/// let mut iter = bm.iter();
327+
/// assert_eq!(iter.next_range_back(), Some(4..=5));
328+
/// assert_eq!(iter.next_back(), Some(2));
329+
/// assert_eq!(iter.next_range_back(), Some(1..=1));
330+
/// ```
331+
pub fn next_range_back(&mut self) -> Option<core::ops::RangeInclusive<u32>> {
332+
next_range_back_impl(&mut self.front, &mut self.containers, &mut self.back)
333+
}
200334
}
201335

202336
impl IntoIter {
@@ -245,6 +379,46 @@ impl IntoIter {
245379
pub fn advance_back_to(&mut self, n: u32) {
246380
advance_back_to_impl(n, &mut self.front, &mut self.containers, &mut self.back);
247381
}
382+
383+
/// Returns the range of consecutive set bits from the current position to the end of the current run
384+
///
385+
/// After this call, the iterator will be positioned at the first item after the returned range.
386+
/// Returns `None` if the iterator is exhausted.
387+
///
388+
/// # Examples
389+
///
390+
/// ```rust
391+
/// use roaring::RoaringBitmap;
392+
///
393+
/// let bm = RoaringBitmap::from([1, 2, 4, 5]);
394+
/// let mut iter = bm.into_iter();
395+
/// assert_eq!(iter.next_range(), Some(1..=2));
396+
/// assert_eq!(iter.next(), Some(4));
397+
/// assert_eq!(iter.next_range(), Some(5..=5));
398+
/// ```
399+
pub fn next_range(&mut self) -> Option<core::ops::RangeInclusive<u32>> {
400+
next_range_impl(&mut self.front, &mut self.containers, &mut self.back)
401+
}
402+
403+
/// Returns the range of consecutive set bits from the start of the current run to the current back position
404+
///
405+
/// After this call, the back of the iterator will be positioned at the last item before the returned range.
406+
/// Returns `None` if the iterator is exhausted.
407+
///
408+
/// # Examples
409+
///
410+
/// ```rust
411+
/// use roaring::RoaringBitmap;
412+
///
413+
/// let bm = RoaringBitmap::from([1, 2, 4, 5]);
414+
/// let mut iter = bm.into_iter();
415+
/// assert_eq!(iter.next_range_back(), Some(4..=5));
416+
/// assert_eq!(iter.next_back(), Some(2));
417+
/// assert_eq!(iter.next_range_back(), Some(1..=1));
418+
/// ```
419+
pub fn next_range_back(&mut self) -> Option<core::ops::RangeInclusive<u32>> {
420+
next_range_back_impl(&mut self.front, &mut self.containers, &mut self.back)
421+
}
248422
}
249423

250424
fn size_hint_impl(

‎roaring/src/bitmap/store/array_store/mod.rs‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,38 @@ pub(crate) struct ArrayStore {
2525
vec: Vec<u16>,
2626
}
2727

28+
/// Return the first contiguous range of elements in a sorted slice.
29+
pub(crate) fn first_contiguous_range_len(slice: &[u16]) -> usize {
30+
let [first, rest @ ..] = slice else {
31+
// Explicitly empty range
32+
return 0;
33+
};
34+
let len = rest.partition_point(|item| {
35+
let item_ptr = core::ptr::addr_of!(*item);
36+
// SAFETY: `item` is guaranteed to be in bounds of `slice`.
37+
let elem_distance = usize::try_from(unsafe { item_ptr.offset_from(first) }).unwrap();
38+
let value_distance = item.checked_sub(*first).expect("array must be sorted");
39+
elem_distance == usize::from(value_distance)
40+
});
41+
len + 1 // +1 for the first element
42+
}
43+
44+
/// Return the first contiguous range of elements in a sorted slice.
45+
pub(crate) fn last_contiguous_range_len(slice: &[u16]) -> usize {
46+
let [rest @ .., last] = slice else {
47+
// Explicitly empty range
48+
return 0;
49+
};
50+
let last_ptr = core::ptr::addr_of!(*last);
51+
let len_from_start = rest.partition_point(|item| {
52+
// SAFETY: `item` is guaranteed to be in bounds of `slice`.
53+
let elem_distance = usize::try_from(unsafe { last_ptr.offset_from(item) }).unwrap();
54+
let value_distance = last.checked_sub(*item).expect("array must be sorted");
55+
elem_distance != usize::from(value_distance)
56+
});
57+
slice.len() - len_from_start
58+
}
59+
2860
impl ArrayStore {
2961
pub fn new() -> ArrayStore {
3062
ArrayStore { vec: vec![] }

0 commit comments

Comments
 (0)