Skip to content

Commit 9f69285

Browse files
KerollmopsKerollmops
authored
Merge pull request #253 from shimatar0/remove-first-last
Add `remove_front` and `remove_back` method
2 parents 6aabe69 + d63dd69 commit 9f69285

File tree

5 files changed

+348
-0
lines changed

5 files changed

+348
-0
lines changed

‎src/bitmap/container.rs‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,36 @@ impl Container {
9494
result
9595
}
9696

97+
pub fn remove_front(&mut self, n: u64) {
98+
match &self.store {
99+
Store::Bitmap(bits) => {
100+
if bits.len() - n <= ARRAY_LIMIT {
101+
let mut replace_array = Vec::with_capacity((bits.len() - n) as usize);
102+
replace_array.extend(bits.iter().skip(n as usize));
103+
self.store = Store::Array(store::ArrayStore::from_vec_unchecked(replace_array));
104+
} else {
105+
self.store.remove_front(n)
106+
}
107+
}
108+
Store::Array(_) => self.store.remove_front(n),
109+
};
110+
}
111+
112+
pub fn remove_back(&mut self, n: u64) {
113+
match &self.store {
114+
Store::Bitmap(bits) => {
115+
if bits.len() - n <= ARRAY_LIMIT {
116+
let mut replace_array = Vec::with_capacity((bits.len() - n) as usize);
117+
replace_array.extend(bits.iter().take((bits.len() - n) as usize));
118+
self.store = Store::Array(store::ArrayStore::from_vec_unchecked(replace_array));
119+
} else {
120+
self.store.remove_back(n)
121+
}
122+
}
123+
Store::Array(_) => self.store.remove_back(n),
124+
};
125+
}
126+
97127
pub fn contains(&self, index: u16) -> bool {
98128
self.store.contains(index)
99129
}

‎src/bitmap/inherent.rs‎

Lines changed: 195 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -578,6 +578,76 @@ impl RoaringBitmap {
578578

579579
None
580580
}
581+
582+
/// Removes the specified number of elements from the top.
583+
///
584+
/// # Examples
585+
///
586+
/// ```rust
587+
/// use roaring::RoaringBitmap;
588+
///
589+
/// let mut rb = RoaringBitmap::from_iter([1, 5, 7, 9]);
590+
/// rb.remove_front(2);
591+
/// assert_eq!(rb, RoaringBitmap::from_iter([7, 9]));
592+
///
593+
/// let mut rb = RoaringBitmap::from_iter([1, 3, 7, 9]);
594+
/// rb.remove_front(2);
595+
/// assert_eq!(rb, RoaringBitmap::from_iter([7, 9]));
596+
pub fn remove_front(&mut self, mut n: u64) {
597+
// remove containers up to the front of the target
598+
let position = self.containers.iter().position(|container| {
599+
let container_len = container.len();
600+
if container_len <= n {
601+
n -= container_len;
602+
false
603+
} else {
604+
true
605+
}
606+
});
607+
let position = position.unwrap_or(self.containers.len());
608+
if position > 0 {
609+
self.containers.drain(..position);
610+
}
611+
// remove data in containers if there are still targets for deletion
612+
if n > 0 && !self.containers.is_empty() {
613+
// container immediately before should have been deleted, so the target is 0 index
614+
self.containers[0].remove_front(n);
615+
}
616+
}
617+
618+
/// Removes the specified number of elements from the tail.
619+
///
620+
/// # Examples
621+
///
622+
/// ```rust
623+
/// use roaring::RoaringBitmap;
624+
///
625+
/// let mut rb = RoaringBitmap::from_iter([1, 5, 7, 9]);
626+
/// rb.remove_back(2);
627+
/// assert_eq!(rb, RoaringBitmap::from_iter([1, 5]));
628+
/// rb.remove_back(1);
629+
/// assert_eq!(rb, RoaringBitmap::from_iter([1]));
630+
pub fn remove_back(&mut self, mut n: u64) {
631+
// remove containers up to the back of the target
632+
let position = self.containers.iter().rposition(|container| {
633+
let container_len = container.len();
634+
if container_len <= n {
635+
n -= container_len;
636+
false
637+
} else {
638+
true
639+
}
640+
});
641+
// It is checked at the beginning of the function, so it is usually never an Err
642+
if let Some(position) = position {
643+
self.containers.drain(position + 1..);
644+
if n > 0 && !self.containers.is_empty() {
645+
self.containers[position].remove_back(n);
646+
}
647+
} else {
648+
self.containers.clear();
649+
}
650+
}
581651
}
582652

583653
impl Default for RoaringBitmap {
@@ -728,4 +798,129 @@ mod tests {
728798
assert_eq!(bitmap.containers.len(), 1);
729799
assert_eq!(bitmap.containers[0].key, 1);
730800
}
801+
802+
#[test]
803+
fn remove_front_for_vec() {
804+
let mut bitmap = RoaringBitmap::from_iter([1, 2, 3, 7, 9, 11]);
805+
bitmap.remove_front(3);
806+
assert_eq!(bitmap.len(), 3);
807+
assert_eq!(bitmap, RoaringBitmap::from_iter([7, 9, 11]));
808+
809+
bitmap = RoaringBitmap::from_iter([1, 2, 5, 7, 9, 11]);
810+
bitmap.remove_front(3);
811+
assert_eq!(bitmap.len(), 3);
812+
assert_eq!(bitmap, RoaringBitmap::from_iter([7, 9, 11]));
813+
814+
bitmap = RoaringBitmap::from_iter([1, 3]);
815+
bitmap.remove_front(2);
816+
assert_eq!(bitmap.len(), 0);
817+
818+
bitmap = RoaringBitmap::from_iter([1, 2, 3, 7, 9, 11]);
819+
bitmap.remove_front(0);
820+
assert_eq!(bitmap.len(), 6);
821+
assert_eq!(bitmap, RoaringBitmap::from_iter([1, 2, 3, 7, 9, 11]));
822+
823+
bitmap = RoaringBitmap::new();
824+
bitmap.insert_range(0..(1_u32 << 16) + 5);
825+
bitmap.remove_front(65537);
826+
assert_eq!(bitmap.len(), 4);
827+
assert_eq!(bitmap, RoaringBitmap::from_iter([65537, 65538, 65539, 65540]));
828+
829+
bitmap = RoaringBitmap::from_iter([1, 2, 5, 7, 9, 11]);
830+
bitmap.remove_front(7);
831+
assert_eq!(bitmap, RoaringBitmap::default());
832+
}
833+
834+
#[test]
835+
fn remove_front_for_bit() {
836+
let mut bitmap = RoaringBitmap::new();
837+
bitmap.insert_range(0..4098);
838+
bitmap.remove_front(4095);
839+
assert_eq!(bitmap.len(), 3);
840+
// removed bit to vec
841+
assert_eq!(bitmap, RoaringBitmap::from_iter([4095, 4096, 4097]));
842+
843+
bitmap = RoaringBitmap::new();
844+
bitmap.insert_range(0..6000);
845+
bitmap.remove_front(999);
846+
assert_eq!(bitmap.len(), 5001);
847+
848+
bitmap = RoaringBitmap::new();
849+
bitmap.insert_range(0..8000);
850+
bitmap.remove_front(10);
851+
assert_eq!(bitmap.len(), 7990);
852+
853+
bitmap = RoaringBitmap::new();
854+
bitmap.insert_range(0..200000);
855+
bitmap.remove_front(2000);
856+
assert_eq!(bitmap.len(), 198000);
857+
assert_eq!(bitmap, RoaringBitmap::from_iter(2000..200000));
858+
859+
bitmap = RoaringBitmap::new();
860+
bitmap.insert_range(0..2);
861+
bitmap.insert_range(4..7);
862+
bitmap.insert_range(1000..6000);
863+
bitmap.remove_front(30);
864+
assert_eq!(bitmap.len(), 4975);
865+
866+
bitmap = RoaringBitmap::new();
867+
bitmap.insert_range(0..65535);
868+
bitmap.remove_front(0);
869+
assert_eq!(bitmap.len(), 65535);
870+
}
871+
872+
#[test]
873+
fn remove_back_for_bit() {
874+
let mut bitmap = RoaringBitmap::new();
875+
bitmap.insert_range(0..5000);
876+
bitmap.remove_back(1000);
877+
assert_eq!(bitmap.len(), 4000);
878+
879+
bitmap = RoaringBitmap::new();
880+
bitmap.insert_range(0..6000);
881+
bitmap.remove_back(1000);
882+
assert_eq!(bitmap.len(), 5000);
883+
884+
bitmap = RoaringBitmap::new();
885+
bitmap.insert_range(0..200000);
886+
bitmap.remove_back(196000);
887+
assert_eq!(bitmap.len(), 4000);
888+
889+
bitmap = RoaringBitmap::new();
890+
bitmap.insert_range(0..200000);
891+
bitmap.remove_back(2000);
892+
assert_eq!(bitmap.len(), 198000);
893+
assert_eq!(bitmap, RoaringBitmap::from_iter(0..198000));
894+
895+
bitmap = RoaringBitmap::new();
896+
bitmap.insert_range(0..65535);
897+
bitmap.remove_back(0);
898+
assert_eq!(bitmap.len(), 65535);
899+
}
900+
901+
#[test]
902+
fn remove_back_for_vec() {
903+
let mut bitmap = RoaringBitmap::from_iter([1, 2, 3, 7, 9, 11]);
904+
bitmap.remove_back(2);
905+
assert_eq!(bitmap, RoaringBitmap::from_iter([1, 2, 3, 7]));
906+
907+
bitmap = RoaringBitmap::from_iter([1, 2, 3, 7, 9, 11]);
908+
bitmap.remove_back(6);
909+
assert_eq!(bitmap.len(), 0);
910+
911+
bitmap = RoaringBitmap::from_iter([1, 2, 3, 7, 9, 11]);
912+
bitmap.remove_back(0);
913+
assert_eq!(bitmap.len(), 6);
914+
assert_eq!(bitmap, RoaringBitmap::from_iter([1, 2, 3, 7, 9, 11]));
915+
916+
bitmap = RoaringBitmap::new();
917+
bitmap.insert_range(0..(1_u32 << 16) + 5);
918+
bitmap.remove_back(65537);
919+
assert_eq!(bitmap.len(), 4);
920+
assert_eq!(bitmap, RoaringBitmap::from_iter([0, 1, 2, 3]));
921+
922+
let mut bitmap = RoaringBitmap::from_iter([1, 2, 3]);
923+
bitmap.remove_back(4);
924+
assert_eq!(bitmap, RoaringBitmap::default());
925+
}
731926
}

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,15 @@ impl ArrayStore {
110110
(pos_end - pos_start) as u64
111111
}
112112

113+
pub fn remove_front(&mut self, n: u64) {
114+
self.vec.rotate_left(n as usize);
115+
self.vec.truncate(self.vec.len() - n as usize);
116+
}
117+
118+
pub fn remove_back(&mut self, n: u64) {
119+
self.vec.truncate(self.vec.len() - n as usize);
120+
}
121+
113122
pub fn contains(&self, index: u16) -> bool {
114123
self.vec.binary_search(&index).is_ok()
115124
}
@@ -562,4 +571,18 @@ mod tests {
562571

563572
assert_eq!(into_vec(store), want);
564573
}
574+
575+
#[test]
576+
fn test_bitmap_remove_front() {
577+
let mut store = Store::Array(ArrayStore::from_vec_unchecked(vec![1, 2, 130, 500]));
578+
store.remove_front(3);
579+
assert_eq!(into_vec(store), vec![500]);
580+
}
581+
582+
#[test]
583+
fn test_bitmap_remove_back() {
584+
let mut store = Store::Array(ArrayStore::from_vec_unchecked(vec![1, 2, 130, 500]));
585+
store.remove_back(2);
586+
assert_eq!(into_vec(store), vec![1, 2]);
587+
}
565588
}

‎src/bitmap/store/bitmap_store.rs‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -302,6 +302,57 @@ impl BitmapStore {
302302
pub fn as_array(&self) -> &[u64; BITMAP_LENGTH] {
303303
&self.bits
304304
}
305+
306+
pub fn clear(&mut self) {
307+
self.bits.fill(0);
308+
self.len = 0;
309+
}
310+
311+
/// Set N bits that are currently 1 bit from the lower bit to 0.
312+
pub fn remove_front(&mut self, mut clear_bits: u64) {
313+
if self.len() < clear_bits {
314+
self.clear();
315+
return;
316+
}
317+
self.len -= clear_bits;
318+
for word in self.bits.iter_mut() {
319+
let count = word.count_ones() as u64;
320+
if clear_bits < count {
321+
for _ in 0..clear_bits {
322+
*word = *word & (*word - 1);
323+
}
324+
return;
325+
}
326+
*word = 0;
327+
clear_bits -= count;
328+
if clear_bits == 0 {
329+
return;
330+
}
331+
}
332+
}
333+
334+
/// Set N bits that are currently 1 bit from the lower bit to 0.
335+
pub fn remove_back(&mut self, mut clear_bits: u64) {
336+
if self.len() < clear_bits {
337+
self.clear();
338+
return;
339+
}
340+
self.len -= clear_bits;
341+
for word in self.bits.iter_mut().rev() {
342+
let count = word.count_ones() as u64;
343+
if clear_bits < count {
344+
for _ in 0..clear_bits {
345+
*word &= !(1 << (63 - word.leading_zeros()));
346+
}
347+
return;
348+
}
349+
*word = 0;
350+
clear_bits -= count;
351+
if clear_bits == 0 {
352+
return;
353+
}
354+
}
355+
}
305356
}
306357

307358
// this can be done in 3 instructions on x86-64 with bmi2 with: tzcnt(pdep(1 << rank, value))
@@ -490,3 +541,38 @@ impl BitXorAssign<&ArrayStore> for BitmapStore {
490541
self.len = len as u64;
491542
}
492543
}
544+
545+
#[cfg(test)]
546+
mod tests {
547+
use super::*;
548+
549+
#[test]
550+
fn test_bitmap_remove_front() {
551+
let mut store = BitmapStore::new();
552+
let range = RangeInclusive::new(1, 3);
553+
store.insert_range(range);
554+
let range_second = RangeInclusive::new(5, 65535);
555+
// store.bits[0] = 0b1111111111111111111111111111111111111111111111111111111111101110
556+
store.insert_range(range_second);
557+
store.remove_front(2);
558+
assert_eq!(
559+
store.bits[0],
560+
0b1111111111111111111111111111111111111111111111111111111111101000
561+
);
562+
}
563+
564+
#[test]
565+
fn test_bitmap_remove_back() {
566+
let mut store = BitmapStore::new();
567+
let range = RangeInclusive::new(1, 3);
568+
store.insert_range(range);
569+
let range_second = RangeInclusive::new(5, 65535);
570+
// store.bits[1023] = 0b1111111111111111111111111111111111111111111111111111111111111111
571+
store.insert_range(range_second);
572+
store.remove_back(2);
573+
assert_eq!(
574+
store.bits[1023],
575+
0b11111111111111111111111111111111111111111111111111111111111111
576+
);
577+
}
578+
}

‎src/bitmap/store/mod.rs‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,20 @@ impl Store {
106106
}
107107
}
108108

109+
pub fn remove_front(&mut self, index: u64) {
110+
match self {
111+
Array(vec) => vec.remove_front(index),
112+
Bitmap(bits) => bits.remove_front(index),
113+
}
114+
}
115+
116+
pub fn remove_back(&mut self, index: u64) {
117+
match self {
118+
Array(vec) => vec.remove_back(index),
119+
Bitmap(bits) => bits.remove_back(index),
120+
}
121+
}
122+
109123
pub fn contains(&self, index: u16) -> bool {
110124
match self {
111125
Array(vec) => vec.contains(index),

0 commit comments

Comments
 (0)