Skip to content

WIP/RFC: shift_remove and friends #558

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 18 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
## [Unreleased]

### Added

- Added `shift_remove`, `shift_remove_entry`, `shift_remove_full`, `shift_remove_index` to `IndexMap`.
- Added `bytes::Buf` and `bytes::BufMut` implementations for `Vec`.
- Added `format` macro.
- Added `String::from_utf16`.
Expand Down
268 changes: 265 additions & 3 deletions src/index_map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ use crate::Vec;
/// ```
pub type FnvIndexMap<K, V, const N: usize> = IndexMap<K, V, BuildHasherDefault<FnvHasher>, N>;

#[derive(Clone, Copy, Eq, PartialEq)]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct HashValue(u16);

impl HashValue {
Expand Down Expand Up @@ -116,6 +116,15 @@ impl Pos {
}
}

impl fmt::Debug for Pos {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Pos")
.field("index", &self.index())
.field("hash", &self.hash().0)
.finish()
}
}

enum Insert<K, V> {
Success(Inserted<V>),
Full((K, V)),
Expand Down Expand Up @@ -302,12 +311,12 @@ where
where
F: FnMut(&mut K, &mut V) -> bool,
{
const INIT: Option<Pos> = None;

self.entries
.retain_mut(|entry| keep(&mut entry.key, &mut entry.value));

if self.entries.len() < self.indices.len() {
const INIT: Option<Pos> = None;

for index in self.indices.iter_mut() {
*index = INIT;
}
Expand Down Expand Up @@ -342,6 +351,24 @@ where
}
}

fn shift_remove_found(&mut self, probe: usize, found: usize) -> (K, V) {
self.indices[probe] = None;
let entry = self.entries.remove(found);

// Re-write the indexes for all elements that moved
for index in self.indices.iter_mut() {
if let Some(pos) = index {
if pos.index() > found {
*index = Some(Pos::new(pos.index() - 1, pos.hash()));
}
}
}

self.backward_shift_after_removal(probe);

(entry.key, entry.value)
}

fn backward_shift_after_removal(&mut self, probe_at_remove: usize) {
// backward shift deletion in self.indices
// after probe, shift all non-ideally placed indices backward
Expand Down Expand Up @@ -1231,6 +1258,159 @@ where
.map(|(probe, found)| self.core.remove_found(probe, found).1)
}

/// Removes the key-value pair at position `index` and returns it.
///
/// Like [`Vec::remove`], the pair is removed by shifting all
/// remaining items. This maintains the remaining elements' relative
/// insertion order, but is a more expensive operation
///
/// Returns `None` if `index` is not in `0..len()`.
///
/// Computes in *O*(n) time (average).
///
/// # Examples
///
/// ```
/// use heapless::index_map::FnvIndexMap;
///
/// let mut map = FnvIndexMap::<_, _, 8>::new();
/// map.insert(3, "a").unwrap();
/// map.insert(2, "b").unwrap();
/// map.insert(1, "c").unwrap();
/// let removed = map.shift_remove_index(1);
/// assert_eq!(removed, Some((2, "b")));
/// assert_eq!(map.len(), 2);
/// assert_eq!(map.shift_remove_index(2), None);
/// assert_eq!(map.shift_remove_index(3), None);
///
/// let mut iter = map.iter();
/// assert_eq!(iter.next(), Some((&3, &"a")));
/// assert_eq!(iter.next(), Some((&1, &"c")));
/// assert_eq!(iter.next(), None);
/// ```
pub fn shift_remove_index(&mut self, index: usize) -> Option<(K, V)> {
if index >= self.len() {
return None;
}
self.find(&self.core.entries[index].key)
.map(|(probe, found)| self.core.shift_remove_found(probe, found))
}

/// Removes the key-value pair equivalent to `key` and returns it and
/// the index it had.
///
/// Like [`Vec::remove`], the pair is removed by shifting all
/// remaining items. This maintains the remaining elements' relative
/// insertion order, but is a more expensive operation.
///
/// Returns `None` if `key` is not in map.
///
/// Computes in *O*(n) time (average).
///
/// # Examples
///
/// ```
/// use heapless::index_map::FnvIndexMap;
///
/// let mut map = FnvIndexMap::<_, _, 8>::new();
/// map.insert(3, "a").unwrap();
/// map.insert(2, "b").unwrap();
/// map.insert(1, "c").unwrap();
/// let removed = map.shift_remove_full(&2);
/// assert_eq!(removed, Some((1, 2, "b")));
/// assert_eq!(map.len(), 2);
/// assert_eq!(map.shift_remove_full(&2), None);
///
/// let mut iter = map.iter();
/// assert_eq!(iter.next(), Some((&3, &"a")));
/// assert_eq!(iter.next(), Some((&1, &"c")));
/// assert_eq!(iter.next(), None);
/// ```
pub fn shift_remove_full<Q>(&mut self, key: &Q) -> Option<(usize, K, V)>
where
K: Borrow<Q>,
Q: ?Sized + Hash + Eq,
{
self.find(key).map(|(probe, found)| {
let (k, v) = self.core.shift_remove_found(probe, found);
(found, k, v)
})
}

/// Removes and returns the key-value pair equivalent to `key`.
///
/// Like [`Vec::remove`], the pair is removed by shifting all
/// remaining items. This maintains the remaining elements' relative
/// insertion order, but is a more expensive operation.
///
/// Returns `None` if `key` is not in map.
///
/// Computes in *O*(n) time (average).
///
/// # Examples
///
/// ```
/// use heapless::index_map::FnvIndexMap;
///
/// let mut map = FnvIndexMap::<_, _, 8>::new();
/// map.insert(3, "a").unwrap();
/// map.insert(2, "b").unwrap();
/// map.insert(1, "c").unwrap();
/// let removed = map.shift_remove_entry(&2);
/// assert_eq!(removed, Some((2, "b")));
/// assert_eq!(map.len(), 2);
/// assert_eq!(map.shift_remove_entry(&2), None);
///
/// let mut iter = map.iter();
/// assert_eq!(iter.next(), Some((&3, &"a")));
/// assert_eq!(iter.next(), Some((&1, &"c")));
/// assert_eq!(iter.next(), None);
/// ```
pub fn shift_remove_entry<Q>(&mut self, key: &Q) -> Option<(K, V)>
where
K: Borrow<Q>,
Q: ?Sized + Hash + Eq,
{
self.shift_remove_full(key).map(|(_idx, k, v)| (k, v))
}

/// Remove the key-value pair equivalent to `key` and return
/// its value.
///
/// Like [`Vec::remove`], the pair is removed by shifting all of the
/// elements that follow it, preserving their relative order.
///
/// Returns `None` if `key` is not in map.
///
/// Computes in *O*(n) time (average).
///
/// # Examples
///
/// ```
/// use heapless::index_map::FnvIndexMap;
///
/// let mut map = FnvIndexMap::<_, _, 8>::new();
/// map.insert(3, "a").unwrap();
/// map.insert(2, "b").unwrap();
/// map.insert(1, "c").unwrap();
/// let removed = map.shift_remove(&2);
/// assert_eq!(removed, Some(("b")));
/// assert_eq!(map.len(), 2);
/// assert_eq!(map.shift_remove(&2), None);
///
/// let mut iter = map.iter();
/// assert_eq!(iter.next(), Some((&3, &"a")));
/// assert_eq!(iter.next(), Some((&1, &"c")));
/// assert_eq!(iter.next(), None);
/// ```
pub fn shift_remove<Q>(&mut self, key: &Q) -> Option<V>
where
K: Borrow<Q>,
Q: ?Sized + Hash + Eq,
{
self.shift_remove_full(key).map(|(_idx, _k, v)| v)
}

/// Retains only the elements specified by the predicate.
///
/// In other words, remove all pairs `(k, v)` for which `f(&k, &mut v)` returns `false`.
Expand Down Expand Up @@ -1281,6 +1461,42 @@ where
let h = hash_with(key, &self.build_hasher);
self.core.find(h, key)
}

#[cfg(test)]
fn assert_internally_consistent(&mut self) {
// For every entry
for (index, bucket) in self.core.entries.iter().enumerate() {
let found = self.find(&bucket.key);
// We must be able to find it
assert!(found.is_some(), "Entry couldn't be found");
let (probe, found_index) = found.unwrap();
// In the right place
assert_eq!(found_index, index, "Found a different entry");
// For the right reason
assert_eq!(
self.core.indices.get(probe),
Some(&Some(Pos::new(index, bucket.hash))),
"Found hash mismatch"
);
}

// For every non-empty index
let mut used_index_count = 0;
for pos in self.core.indices.iter().filter_map(|p| p.as_ref()) {
let entry = self.core.entries.get(pos.index());
// It must point to an entry
assert!(entry.is_some(), "Dangling reference");
// With a matching hash
assert_eq!(pos.hash(), entry.unwrap().hash, "Index hash mismatch");
used_index_count += 1;
}

assert_eq!(
self.core.entries.len(),
used_index_count,
"Entry/index count mismatch"
);
}
}

impl<K, Q, V, S, const N: usize> ops::Index<&Q> for IndexMap<K, V, S, N>
Expand Down Expand Up @@ -1850,22 +2066,67 @@ mod tests {
assert_eq!(MAP_SLOTS - 1, src.len());
}

#[test]
fn shift_remove_full() {
const REMOVED_KEY: usize = 7;
let mut map = almost_filled_map();
map.assert_internally_consistent();
assert_eq!(
map.shift_remove_full(&REMOVED_KEY),
Some((REMOVED_KEY - 1, REMOVED_KEY, REMOVED_KEY))
);
// Verify all other elements can still be looked up after removing the entry
for x in 1..MAP_SLOTS - 1 {
if x == REMOVED_KEY {
continue;
}
assert!(map.contains_key(&x));
}
// Make sure we can't find any entry for the one we removed
assert_eq!(map.get(&REMOVED_KEY), None);
map.assert_internally_consistent();
}

#[test]
fn shift_remove_index() {
const REMOVED_KEY: usize = 7;
let mut map = almost_filled_map();
map.assert_internally_consistent();
assert_eq!(
map.shift_remove_index(REMOVED_KEY - 1),
Some((REMOVED_KEY, REMOVED_KEY))
);
// Verify all other elements can still be looked up after removing the entry
for x in 1..MAP_SLOTS - 1 {
if x == REMOVED_KEY {
continue;
}
assert!(map.contains_key(&x));
}
// Make sure we can't find any entry for the one we removed
assert_eq!(map.get(&REMOVED_KEY), None);
map.assert_internally_consistent();
}

#[test]
fn retain() {
let mut none = almost_filled_map();
none.retain(|_, _| false);
assert!(none.is_empty());
none.assert_internally_consistent();

let mut all = almost_filled_map();
all.retain(|_, _| true);
assert_eq!(all.len(), MAP_SLOTS - 1);
all.assert_internally_consistent();

let mut even = almost_filled_map();
even.retain(|_, &mut v| v % 2 == 0);
assert_eq!(even.len(), (MAP_SLOTS - 1) / 2);
for &v in even.values() {
assert_eq!(v % 2, 0);
}
even.assert_internally_consistent();

let mut odd = almost_filled_map();
odd.retain(|_, &mut v| v % 2 != 0);
Expand All @@ -1875,6 +2136,7 @@ mod tests {
}
assert_eq!(odd.insert(2, 2), Ok(None));
assert_eq!(odd.len(), (MAP_SLOTS / 2) + 1);
odd.assert_internally_consistent();
}

#[test]
Expand Down