From 63788568c0a9e0f5959721131aec75799a40f647 Mon Sep 17 00:00:00 2001 From: Josh Stone Date: Wed, 30 Sep 2020 12:01:45 -0700 Subject: [PATCH] Add more Vec/slice-like methods to maps and sets ```rust impl IndexMap { pub fn truncate(&mut self, len: usize); pub fn split_off(&mut self, at: usize) -> Self where S: Clone; pub fn first(&self) -> Option<(&K, &V)>; pub fn first_mut(&mut self) -> Option<(&mut K, &mut V)>; pub fn last(&self) -> Option<(&K, &V)>; pub fn last_mut(&mut self) -> Option<(&mut K, &mut V)>; pub fn swap_indices(&mut self, a: usize, b: usize); } impl IndexSet { pub fn truncate(&mut self, len: usize); pub fn split_off(&mut self, at: usize) -> Self where S: Clone; pub fn first(&self) -> Option<&T>; pub fn last(&self) -> Option<&T>; pub fn swap_indices(&mut self, a: usize, b: usize); } ``` I would prefer `&K` instead of `&mut K` in `IndexMap::first_mut` and `last_mut`, but this is consistent with the existing `get_index_mut`. --- src/map.rs | 47 +++++++++++++++++++++++++++++++++++++++++++++ src/map/core.rs | 19 ++++++++++++++++++ src/map/core/raw.rs | 23 ++++++++++++++++++++++ src/set.rs | 44 +++++++++++++++++++++++++++++++++++++++--- 4 files changed, 130 insertions(+), 3 deletions(-) diff --git a/src/map.rs b/src/map.rs index 6829ef94..467fee81 100644 --- a/src/map.rs +++ b/src/map.rs @@ -252,6 +252,13 @@ impl IndexMap { self.core.clear(); } + /// Shortens the map, keeping the first `len` elements and dropping the rest. + /// + /// If `len` is greater than the map's current length, this has no effect. + pub fn truncate(&mut self, len: usize) { + self.core.truncate(len); + } + /// Clears the `IndexMap` in the given index range, returning those /// key-value pairs as a drain iterator. /// @@ -273,6 +280,23 @@ impl IndexMap { iter: self.core.drain(range), } } + + /// Splits the collection into two at the given index. + /// + /// Returns a newly allocated map containing the elements in the range + /// `[at, len)`. After the call, the original map will be left containing + /// the elements `[0, at)` with its previous capacity unchanged. + /// + /// ***Panics*** if `at > len`. + pub fn split_off(&mut self, at: usize) -> Self + where + S: Clone, + { + Self { + core: self.core.split_off(at), + hash_builder: self.hash_builder.clone(), + } + } } impl IndexMap @@ -693,6 +717,22 @@ impl IndexMap { self.as_entries_mut().get_mut(index).map(Bucket::muts) } + pub fn first(&self) -> Option<(&K, &V)> { + self.as_entries().first().map(Bucket::refs) + } + + pub fn first_mut(&mut self) -> Option<(&mut K, &mut V)> { + self.as_entries_mut().first_mut().map(Bucket::muts) + } + + pub fn last(&self) -> Option<(&K, &V)> { + self.as_entries().last().map(Bucket::refs) + } + + pub fn last_mut(&mut self) -> Option<(&mut K, &mut V)> { + self.as_entries_mut().last_mut().map(Bucket::muts) + } + /// Remove the key-value pair by index /// /// Valid indices are *0 <= index < self.len()* @@ -718,6 +758,13 @@ impl IndexMap { pub fn shift_remove_index(&mut self, index: usize) -> Option<(K, V)> { self.core.shift_remove_index(index) } + + /// Swaps the position of two key-value pairs in the map. + /// + /// ***Panics*** if `a` or `b` are out of bounds. + pub fn swap_indices(&mut self, a: usize, b: usize) { + self.core.swap_indices(a, b) + } } /// An iterator over the keys of a `IndexMap`. diff --git a/src/map/core.rs b/src/map/core.rs index 93876b2e..02e99e07 100644 --- a/src/map/core.rs +++ b/src/map/core.rs @@ -150,6 +150,13 @@ impl IndexMapCore { self.entries.clear(); } + pub(crate) fn truncate(&mut self, len: usize) { + if len < self.len() { + self.erase_indices(len, self.entries.len()); + self.entries.truncate(len); + } + } + pub(crate) fn drain(&mut self, range: R) -> Drain<'_, Bucket> where R: RangeBounds, @@ -159,6 +166,18 @@ impl IndexMapCore { self.entries.drain(range) } + pub(crate) fn split_off(&mut self, at: usize) -> Self { + assert!(at <= self.entries.len()); + self.erase_indices(at, self.entries.len()); + let entries = self.entries.split_off(at); + + let mut indices = RawTable::with_capacity(entries.len()); + for (i, entry) in enumerate(&entries) { + indices.insert_no_grow(entry.hash.get(), i); + } + Self { indices, entries } + } + /// Reserve capacity for `additional` more key-value pairs. pub(crate) fn reserve(&mut self, additional: usize) { self.indices.reserve(additional, get_hash(&self.entries)); diff --git a/src/map/core/raw.rs b/src/map/core/raw.rs index 82d92a48..3388f123 100644 --- a/src/map/core/raw.rs +++ b/src/map/core/raw.rs @@ -61,6 +61,29 @@ impl IndexMapCore { // only the item references that are appropriately bound to `&mut self`. unsafe { self.indices.iter().map(|bucket| bucket.as_mut()) } } + + /// Return the raw bucket for the given index + fn find_index(&self, index: usize) -> RawBucket { + // We'll get a "nice" bounds-check from indexing `self.entries`, + // and then we expect to find it in the table as well. + let hash = self.entries[index].hash.get(); + self.indices + .find(hash, move |&i| i == index) + .expect("index not found") + } + + pub(crate) fn swap_indices(&mut self, a: usize, b: usize) { + // SAFETY: Can't take two `get_mut` references from one table, so we + // must use raw buckets to do the swap. This is still safe because we + // are locally sure they won't dangle, and we write them individually. + unsafe { + let raw_bucket_a = self.find_index(a); + let raw_bucket_b = self.find_index(b); + raw_bucket_a.write(b); + raw_bucket_b.write(a); + } + self.entries.swap(a, b); + } } /// A view into an occupied entry in a `IndexMap`. diff --git a/src/set.rs b/src/set.rs index 99493891..af1ef451 100644 --- a/src/set.rs +++ b/src/set.rs @@ -200,6 +200,13 @@ impl IndexSet { self.map.clear(); } + /// Shortens the set, keeping the first `len` elements and dropping the rest. + /// + /// If `len` is greater than the set's current length, this has no effect. + pub fn truncate(&mut self, len: usize) { + self.map.truncate(len); + } + /// Clears the `IndexSet` in the given index range, returning those values /// as a drain iterator. /// @@ -221,6 +228,22 @@ impl IndexSet { iter: self.map.drain(range).iter, } } + + /// Splits the collection into two at the given index. + /// + /// Returns a newly allocated set containing the elements in the range + /// `[at, len)`. After the call, the original set will be left containing + /// the elements `[0, at)` with its previous capacity unchanged. + /// + /// ***Panics*** if `at > len`. + pub fn split_off(&mut self, at: usize) -> Self + where + S: Clone, + { + Self { + map: self.map.split_off(at), + } + } } impl IndexSet @@ -576,10 +599,18 @@ impl IndexSet { /// /// Computes in **O(1)** time. pub fn get_index(&self, index: usize) -> Option<&T> { - self.map.get_index(index).map(|(x, &())| x) + self.as_entries().get(index).map(Bucket::key_ref) + } + + pub fn first(&self) -> Option<&T> { + self.as_entries().first().map(Bucket::key_ref) + } + + pub fn last(&self) -> Option<&T> { + self.as_entries().last().map(Bucket::key_ref) } - /// Remove the key-value pair by index + /// Remove the value by index /// /// Valid indices are *0 <= index < self.len()* /// @@ -592,7 +623,7 @@ impl IndexSet { self.map.swap_remove_index(index).map(|(x, ())| x) } - /// Remove the key-value pair by index + /// Remove the value by index /// /// Valid indices are *0 <= index < self.len()* /// @@ -604,6 +635,13 @@ impl IndexSet { pub fn shift_remove_index(&mut self, index: usize) -> Option { self.map.shift_remove_index(index).map(|(x, ())| x) } + + /// Swaps the position of two values in the set. + /// + /// ***Panics*** if `a` or `b` are out of bounds. + pub fn swap_indices(&mut self, a: usize, b: usize) { + self.map.swap_indices(a, b) + } } /// Access `IndexSet` values at indexed positions.