109 lines
3.6 KiB
Rust
109 lines
3.6 KiB
Rust
#![allow(unsafe_code)]
|
|
|
|
use super::{Bucket, Core};
|
|
use crate::util::simplify_range;
|
|
|
|
use core::ops::RangeBounds;
|
|
|
|
impl<K, V> Core<K, V> {
|
|
#[track_caller]
|
|
pub(crate) fn extract<R>(&mut self, range: R) -> ExtractCore<'_, K, V>
|
|
where
|
|
R: RangeBounds<usize>,
|
|
{
|
|
let range = simplify_range(range, self.entries.len());
|
|
|
|
// SAFETY: We must have consistent lengths to start, so that's a hard assertion.
|
|
// Then the worst `set_len` can do is leak items if `ExtractCore` doesn't drop.
|
|
assert_eq!(self.entries.len(), self.indices.len());
|
|
unsafe {
|
|
self.entries.set_len(range.start);
|
|
}
|
|
ExtractCore {
|
|
map: self,
|
|
new_len: range.start,
|
|
current: range.start,
|
|
end: range.end,
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(crate) struct ExtractCore<'a, K, V> {
|
|
map: &'a mut Core<K, V>,
|
|
new_len: usize,
|
|
current: usize,
|
|
end: usize,
|
|
}
|
|
|
|
impl<K, V> Drop for ExtractCore<'_, K, V> {
|
|
fn drop(&mut self) {
|
|
let old_len = self.map.indices.len();
|
|
let mut new_len = self.new_len;
|
|
|
|
debug_assert!(new_len <= self.current);
|
|
debug_assert!(self.current <= self.end);
|
|
debug_assert!(self.current <= old_len);
|
|
debug_assert!(old_len <= self.map.entries.capacity());
|
|
|
|
// SAFETY: We assume `new_len` and `current` were correctly maintained by the iterator.
|
|
// So `entries[new_len..current]` were extracted, but the rest before and after are valid.
|
|
unsafe {
|
|
if new_len == self.current {
|
|
// Nothing was extracted, so any remaining items can be left in place.
|
|
new_len = old_len;
|
|
} else if self.current < old_len {
|
|
// Need to shift the remaining items down.
|
|
let tail_len = old_len - self.current;
|
|
let base = self.map.entries.as_mut_ptr();
|
|
let src = base.add(self.current);
|
|
let dest = base.add(new_len);
|
|
src.copy_to(dest, tail_len);
|
|
new_len += tail_len;
|
|
}
|
|
self.map.entries.set_len(new_len);
|
|
}
|
|
|
|
if new_len != old_len {
|
|
// We don't keep track of *which* items were extracted, so reindex everything.
|
|
self.map.rebuild_hash_table();
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<K, V> ExtractCore<'_, K, V> {
|
|
pub(crate) fn extract_if<F>(&mut self, mut pred: F) -> Option<Bucket<K, V>>
|
|
where
|
|
F: FnMut(&mut Bucket<K, V>) -> bool,
|
|
{
|
|
debug_assert!(self.end <= self.map.entries.capacity());
|
|
|
|
let base = self.map.entries.as_mut_ptr();
|
|
while self.current < self.end {
|
|
// SAFETY: We're maintaining both indices within bounds of the original entries, so
|
|
// 0..new_len and current..indices.len() are always valid items for our Drop to keep.
|
|
unsafe {
|
|
let item = base.add(self.current);
|
|
if pred(&mut *item) {
|
|
// Extract it!
|
|
self.current += 1;
|
|
return Some(item.read());
|
|
} else {
|
|
// Keep it, shifting it down if needed.
|
|
if self.new_len != self.current {
|
|
debug_assert!(self.new_len < self.current);
|
|
let dest = base.add(self.new_len);
|
|
item.copy_to_nonoverlapping(dest, 1);
|
|
}
|
|
self.current += 1;
|
|
self.new_len += 1;
|
|
}
|
|
}
|
|
}
|
|
None
|
|
}
|
|
|
|
pub(crate) fn remaining(&self) -> usize {
|
|
self.end - self.current
|
|
}
|
|
}
|