Skip to content

Commit

Permalink
Merge pull request #253 from shimatar0/remove-first-last
Browse files Browse the repository at this point in the history 
Add `remove_front` and `remove_back` method
  • Loading branch information
@Kerollmops
Kerollmops authored Jul 17, 2023
2 parents 6aabe69 + d63dd69 commit 9f69285
Show file tree
Hide file tree
Showing 5 changed files with 348 additions and 0 deletions.
30 changes: 30 additions & 0 deletions src/bitmap/container.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,36 @@ impl Container {
result
}

pub fn remove_front(&mut self, n: u64) {
match &self.store {
Store::Bitmap(bits) => {
if bits.len() - n <= ARRAY_LIMIT {
let mut replace_array = Vec::with_capacity((bits.len() - n) as usize);
replace_array.extend(bits.iter().skip(n as usize));
self.store = Store::Array(store::ArrayStore::from_vec_unchecked(replace_array));
} else {
self.store.remove_front(n)
}
}
Store::Array(_) => self.store.remove_front(n),
};
}

pub fn remove_back(&mut self, n: u64) {
match &self.store {
Store::Bitmap(bits) => {
if bits.len() - n <= ARRAY_LIMIT {
let mut replace_array = Vec::with_capacity((bits.len() - n) as usize);
replace_array.extend(bits.iter().take((bits.len() - n) as usize));
self.store = Store::Array(store::ArrayStore::from_vec_unchecked(replace_array));
} else {
self.store.remove_back(n)
}
}
Store::Array(_) => self.store.remove_back(n),
};
}

pub fn contains(&self, index: u16) -> bool {
self.store.contains(index)
}
Expand Down
195 changes: 195 additions & 0 deletions src/bitmap/inherent.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -578,6 +578,76 @@ impl RoaringBitmap {

None
}

/// Removes the specified number of elements from the top.
///
/// # Examples
///
/// ```rust
/// use roaring::RoaringBitmap;
///
/// let mut rb = RoaringBitmap::from_iter([1, 5, 7, 9]);
/// rb.remove_front(2);
/// assert_eq!(rb, RoaringBitmap::from_iter([7, 9]));
///
/// let mut rb = RoaringBitmap::from_iter([1, 3, 7, 9]);
/// rb.remove_front(2);
/// assert_eq!(rb, RoaringBitmap::from_iter([7, 9]));
pub fn remove_front(&mut self, mut n: u64) {
// remove containers up to the front of the target
let position = self.containers.iter().position(|container| {
let container_len = container.len();
if container_len <= n {
n -= container_len;
false
} else {
true
}
});
let position = position.unwrap_or(self.containers.len());
if position > 0 {
self.containers.drain(..position);
}
// remove data in containers if there are still targets for deletion
if n > 0 && !self.containers.is_empty() {
// container immediately before should have been deleted, so the target is 0 index
self.containers[0].remove_front(n);
}
}

/// Removes the specified number of elements from the tail.
///
/// # Examples
///
/// ```rust
/// use roaring::RoaringBitmap;
///
/// let mut rb = RoaringBitmap::from_iter([1, 5, 7, 9]);
/// rb.remove_back(2);
/// assert_eq!(rb, RoaringBitmap::from_iter([1, 5]));
/// rb.remove_back(1);
/// assert_eq!(rb, RoaringBitmap::from_iter([1]));
pub fn remove_back(&mut self, mut n: u64) {
// remove containers up to the back of the target
let position = self.containers.iter().rposition(|container| {
let container_len = container.len();
if container_len <= n {
n -= container_len;
false
} else {
true
}
});
// It is checked at the beginning of the function, so it is usually never an Err
if let Some(position) = position {
self.containers.drain(position + 1..);
if n > 0 && !self.containers.is_empty() {
self.containers[position].remove_back(n);
}
} else {
self.containers.clear();
}
}
}

impl Default for RoaringBitmap {
Expand Down Expand Up @@ -728,4 +798,129 @@ mod tests {
assert_eq!(bitmap.containers.len(), 1);
assert_eq!(bitmap.containers[0].key, 1);
}

#[test]
fn remove_front_for_vec() {
let mut bitmap = RoaringBitmap::from_iter([1, 2, 3, 7, 9, 11]);
bitmap.remove_front(3);
assert_eq!(bitmap.len(), 3);
assert_eq!(bitmap, RoaringBitmap::from_iter([7, 9, 11]));

bitmap = RoaringBitmap::from_iter([1, 2, 5, 7, 9, 11]);
bitmap.remove_front(3);
assert_eq!(bitmap.len(), 3);
assert_eq!(bitmap, RoaringBitmap::from_iter([7, 9, 11]));

bitmap = RoaringBitmap::from_iter([1, 3]);
bitmap.remove_front(2);
assert_eq!(bitmap.len(), 0);

bitmap = RoaringBitmap::from_iter([1, 2, 3, 7, 9, 11]);
bitmap.remove_front(0);
assert_eq!(bitmap.len(), 6);
assert_eq!(bitmap, RoaringBitmap::from_iter([1, 2, 3, 7, 9, 11]));

bitmap = RoaringBitmap::new();
bitmap.insert_range(0..(1_u32 << 16) + 5);
bitmap.remove_front(65537);
assert_eq!(bitmap.len(), 4);
assert_eq!(bitmap, RoaringBitmap::from_iter([65537, 65538, 65539, 65540]));

bitmap = RoaringBitmap::from_iter([1, 2, 5, 7, 9, 11]);
bitmap.remove_front(7);
assert_eq!(bitmap, RoaringBitmap::default());
}

#[test]
fn remove_front_for_bit() {
let mut bitmap = RoaringBitmap::new();
bitmap.insert_range(0..4098);
bitmap.remove_front(4095);
assert_eq!(bitmap.len(), 3);
// removed bit to vec
assert_eq!(bitmap, RoaringBitmap::from_iter([4095, 4096, 4097]));

bitmap = RoaringBitmap::new();
bitmap.insert_range(0..6000);
bitmap.remove_front(999);
assert_eq!(bitmap.len(), 5001);

bitmap = RoaringBitmap::new();
bitmap.insert_range(0..8000);
bitmap.remove_front(10);
assert_eq!(bitmap.len(), 7990);

bitmap = RoaringBitmap::new();
bitmap.insert_range(0..200000);
bitmap.remove_front(2000);
assert_eq!(bitmap.len(), 198000);
assert_eq!(bitmap, RoaringBitmap::from_iter(2000..200000));

bitmap = RoaringBitmap::new();
bitmap.insert_range(0..2);
bitmap.insert_range(4..7);
bitmap.insert_range(1000..6000);
bitmap.remove_front(30);
assert_eq!(bitmap.len(), 4975);

bitmap = RoaringBitmap::new();
bitmap.insert_range(0..65535);
bitmap.remove_front(0);
assert_eq!(bitmap.len(), 65535);
}

#[test]
fn remove_back_for_bit() {
let mut bitmap = RoaringBitmap::new();
bitmap.insert_range(0..5000);
bitmap.remove_back(1000);
assert_eq!(bitmap.len(), 4000);

bitmap = RoaringBitmap::new();
bitmap.insert_range(0..6000);
bitmap.remove_back(1000);
assert_eq!(bitmap.len(), 5000);

bitmap = RoaringBitmap::new();
bitmap.insert_range(0..200000);
bitmap.remove_back(196000);
assert_eq!(bitmap.len(), 4000);

bitmap = RoaringBitmap::new();
bitmap.insert_range(0..200000);
bitmap.remove_back(2000);
assert_eq!(bitmap.len(), 198000);
assert_eq!(bitmap, RoaringBitmap::from_iter(0..198000));

bitmap = RoaringBitmap::new();
bitmap.insert_range(0..65535);
bitmap.remove_back(0);
assert_eq!(bitmap.len(), 65535);
}

#[test]
fn remove_back_for_vec() {
let mut bitmap = RoaringBitmap::from_iter([1, 2, 3, 7, 9, 11]);
bitmap.remove_back(2);
assert_eq!(bitmap, RoaringBitmap::from_iter([1, 2, 3, 7]));

bitmap = RoaringBitmap::from_iter([1, 2, 3, 7, 9, 11]);
bitmap.remove_back(6);
assert_eq!(bitmap.len(), 0);

bitmap = RoaringBitmap::from_iter([1, 2, 3, 7, 9, 11]);
bitmap.remove_back(0);
assert_eq!(bitmap.len(), 6);
assert_eq!(bitmap, RoaringBitmap::from_iter([1, 2, 3, 7, 9, 11]));

bitmap = RoaringBitmap::new();
bitmap.insert_range(0..(1_u32 << 16) + 5);
bitmap.remove_back(65537);
assert_eq!(bitmap.len(), 4);
assert_eq!(bitmap, RoaringBitmap::from_iter([0, 1, 2, 3]));

let mut bitmap = RoaringBitmap::from_iter([1, 2, 3]);
bitmap.remove_back(4);
assert_eq!(bitmap, RoaringBitmap::default());
}
}
23 changes: 23 additions & 0 deletions src/bitmap/store/array_store/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,15 @@ impl ArrayStore {
(pos_end - pos_start) as u64
}

pub fn remove_front(&mut self, n: u64) {
self.vec.rotate_left(n as usize);
self.vec.truncate(self.vec.len() - n as usize);
}

pub fn remove_back(&mut self, n: u64) {
self.vec.truncate(self.vec.len() - n as usize);
}

pub fn contains(&self, index: u16) -> bool {
self.vec.binary_search(&index).is_ok()
}
Expand Down Expand Up @@ -562,4 +571,18 @@ mod tests {

assert_eq!(into_vec(store), want);
}

#[test]
fn test_bitmap_remove_front() {
let mut store = Store::Array(ArrayStore::from_vec_unchecked(vec![1, 2, 130, 500]));
store.remove_front(3);
assert_eq!(into_vec(store), vec![500]);
}

#[test]
fn test_bitmap_remove_back() {
let mut store = Store::Array(ArrayStore::from_vec_unchecked(vec![1, 2, 130, 500]));
store.remove_back(2);
assert_eq!(into_vec(store), vec![1, 2]);
}
}
86 changes: 86 additions & 0 deletions src/bitmap/store/bitmap_store.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -302,6 +302,57 @@ impl BitmapStore {
pub fn as_array(&self) -> &[u64; BITMAP_LENGTH] {
&self.bits
}

pub fn clear(&mut self) {
self.bits.fill(0);
self.len = 0;
}

/// Set N bits that are currently 1 bit from the lower bit to 0.
pub fn remove_front(&mut self, mut clear_bits: u64) {
if self.len() < clear_bits {
self.clear();
return;
}
self.len -= clear_bits;
for word in self.bits.iter_mut() {
let count = word.count_ones() as u64;
if clear_bits < count {
for _ in 0..clear_bits {
*word = *word & (*word - 1);
}
return;
}
*word = 0;
clear_bits -= count;
if clear_bits == 0 {
return;
}
}
}

/// Set N bits that are currently 1 bit from the lower bit to 0.
pub fn remove_back(&mut self, mut clear_bits: u64) {
if self.len() < clear_bits {
self.clear();
return;
}
self.len -= clear_bits;
for word in self.bits.iter_mut().rev() {
let count = word.count_ones() as u64;
if clear_bits < count {
for _ in 0..clear_bits {
*word &= !(1 << (63 - word.leading_zeros()));
}
return;
}
*word = 0;
clear_bits -= count;
if clear_bits == 0 {
return;
}
}
}
}

// this can be done in 3 instructions on x86-64 with bmi2 with: tzcnt(pdep(1 << rank, value))
Expand Down Expand Up @@ -490,3 +541,38 @@ impl BitXorAssign<&ArrayStore> for BitmapStore {
self.len = len as u64;
}
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn test_bitmap_remove_front() {
let mut store = BitmapStore::new();
let range = RangeInclusive::new(1, 3);
store.insert_range(range);
let range_second = RangeInclusive::new(5, 65535);
// store.bits[0] = 0b1111111111111111111111111111111111111111111111111111111111101110
store.insert_range(range_second);
store.remove_front(2);
assert_eq!(
store.bits[0],
0b1111111111111111111111111111111111111111111111111111111111101000
);
}

#[test]
fn test_bitmap_remove_back() {
let mut store = BitmapStore::new();
let range = RangeInclusive::new(1, 3);
store.insert_range(range);
let range_second = RangeInclusive::new(5, 65535);
// store.bits[1023] = 0b1111111111111111111111111111111111111111111111111111111111111111
store.insert_range(range_second);
store.remove_back(2);
assert_eq!(
store.bits[1023],
0b11111111111111111111111111111111111111111111111111111111111111
);
}
}
14 changes: 14 additions & 0 deletions src/bitmap/store/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,20 @@ impl Store {
}
}

pub fn remove_front(&mut self, index: u64) {
match self {
Array(vec) => vec.remove_front(index),
Bitmap(bits) => bits.remove_front(index),
}
}

pub fn remove_back(&mut self, index: u64) {
match self {
Array(vec) => vec.remove_back(index),
Bitmap(bits) => bits.remove_back(index),
}
}

pub fn contains(&self, index: u16) -> bool {
match self {
Array(vec) => vec.contains(index),
Expand Down

0 comments on commit 9f69285

Please sign in to comment.