Skip to content

Commit 6378856

Browse files
committed
Add more Vec/slice-like methods to maps and sets
```rust impl<K, V, S> IndexMap<K, V, S> { 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<T, S> IndexSet<T, S> { 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`.
1 parent 59e2833 commit 6378856

File tree

4 files changed

+130
-3
lines changed

4 files changed

+130
-3
lines changed

src/map.rs

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -252,6 +252,13 @@ impl<K, V, S> IndexMap<K, V, S> {
252252
self.core.clear();
253253
}
254254

255+
/// Shortens the map, keeping the first `len` elements and dropping the rest.
256+
///
257+
/// If `len` is greater than the map's current length, this has no effect.
258+
pub fn truncate(&mut self, len: usize) {
259+
self.core.truncate(len);
260+
}
261+
255262
/// Clears the `IndexMap` in the given index range, returning those
256263
/// key-value pairs as a drain iterator.
257264
///
@@ -273,6 +280,23 @@ impl<K, V, S> IndexMap<K, V, S> {
273280
iter: self.core.drain(range),
274281
}
275282
}
283+
284+
/// Splits the collection into two at the given index.
285+
///
286+
/// Returns a newly allocated map containing the elements in the range
287+
/// `[at, len)`. After the call, the original map will be left containing
288+
/// the elements `[0, at)` with its previous capacity unchanged.
289+
///
290+
/// ***Panics*** if `at > len`.
291+
pub fn split_off(&mut self, at: usize) -> Self
292+
where
293+
S: Clone,
294+
{
295+
Self {
296+
core: self.core.split_off(at),
297+
hash_builder: self.hash_builder.clone(),
298+
}
299+
}
276300
}
277301

278302
impl<K, V, S> IndexMap<K, V, S>
@@ -693,6 +717,22 @@ impl<K, V, S> IndexMap<K, V, S> {
693717
self.as_entries_mut().get_mut(index).map(Bucket::muts)
694718
}
695719

720+
pub fn first(&self) -> Option<(&K, &V)> {
721+
self.as_entries().first().map(Bucket::refs)
722+
}
723+
724+
pub fn first_mut(&mut self) -> Option<(&mut K, &mut V)> {
725+
self.as_entries_mut().first_mut().map(Bucket::muts)
726+
}
727+
728+
pub fn last(&self) -> Option<(&K, &V)> {
729+
self.as_entries().last().map(Bucket::refs)
730+
}
731+
732+
pub fn last_mut(&mut self) -> Option<(&mut K, &mut V)> {
733+
self.as_entries_mut().last_mut().map(Bucket::muts)
734+
}
735+
696736
/// Remove the key-value pair by index
697737
///
698738
/// Valid indices are *0 <= index < self.len()*
@@ -718,6 +758,13 @@ impl<K, V, S> IndexMap<K, V, S> {
718758
pub fn shift_remove_index(&mut self, index: usize) -> Option<(K, V)> {
719759
self.core.shift_remove_index(index)
720760
}
761+
762+
/// Swaps the position of two key-value pairs in the map.
763+
///
764+
/// ***Panics*** if `a` or `b` are out of bounds.
765+
pub fn swap_indices(&mut self, a: usize, b: usize) {
766+
self.core.swap_indices(a, b)
767+
}
721768
}
722769

723770
/// An iterator over the keys of a `IndexMap`.

src/map/core.rs

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,13 @@ impl<K, V> IndexMapCore<K, V> {
150150
self.entries.clear();
151151
}
152152

153+
pub(crate) fn truncate(&mut self, len: usize) {
154+
if len < self.len() {
155+
self.erase_indices(len, self.entries.len());
156+
self.entries.truncate(len);
157+
}
158+
}
159+
153160
pub(crate) fn drain<R>(&mut self, range: R) -> Drain<'_, Bucket<K, V>>
154161
where
155162
R: RangeBounds<usize>,
@@ -159,6 +166,18 @@ impl<K, V> IndexMapCore<K, V> {
159166
self.entries.drain(range)
160167
}
161168

169+
pub(crate) fn split_off(&mut self, at: usize) -> Self {
170+
assert!(at <= self.entries.len());
171+
self.erase_indices(at, self.entries.len());
172+
let entries = self.entries.split_off(at);
173+
174+
let mut indices = RawTable::with_capacity(entries.len());
175+
for (i, entry) in enumerate(&entries) {
176+
indices.insert_no_grow(entry.hash.get(), i);
177+
}
178+
Self { indices, entries }
179+
}
180+
162181
/// Reserve capacity for `additional` more key-value pairs.
163182
pub(crate) fn reserve(&mut self, additional: usize) {
164183
self.indices.reserve(additional, get_hash(&self.entries));

src/map/core/raw.rs

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,29 @@ impl<K, V> IndexMapCore<K, V> {
6161
// only the item references that are appropriately bound to `&mut self`.
6262
unsafe { self.indices.iter().map(|bucket| bucket.as_mut()) }
6363
}
64+
65+
/// Return the raw bucket for the given index
66+
fn find_index(&self, index: usize) -> RawBucket {
67+
// We'll get a "nice" bounds-check from indexing `self.entries`,
68+
// and then we expect to find it in the table as well.
69+
let hash = self.entries[index].hash.get();
70+
self.indices
71+
.find(hash, move |&i| i == index)
72+
.expect("index not found")
73+
}
74+
75+
pub(crate) fn swap_indices(&mut self, a: usize, b: usize) {
76+
// SAFETY: Can't take two `get_mut` references from one table, so we
77+
// must use raw buckets to do the swap. This is still safe because we
78+
// are locally sure they won't dangle, and we write them individually.
79+
unsafe {
80+
let raw_bucket_a = self.find_index(a);
81+
let raw_bucket_b = self.find_index(b);
82+
raw_bucket_a.write(b);
83+
raw_bucket_b.write(a);
84+
}
85+
self.entries.swap(a, b);
86+
}
6487
}
6588

6689
/// A view into an occupied entry in a `IndexMap`.

src/set.rs

Lines changed: 41 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,13 @@ impl<T, S> IndexSet<T, S> {
200200
self.map.clear();
201201
}
202202

203+
/// Shortens the set, keeping the first `len` elements and dropping the rest.
204+
///
205+
/// If `len` is greater than the set's current length, this has no effect.
206+
pub fn truncate(&mut self, len: usize) {
207+
self.map.truncate(len);
208+
}
209+
203210
/// Clears the `IndexSet` in the given index range, returning those values
204211
/// as a drain iterator.
205212
///
@@ -221,6 +228,22 @@ impl<T, S> IndexSet<T, S> {
221228
iter: self.map.drain(range).iter,
222229
}
223230
}
231+
232+
/// Splits the collection into two at the given index.
233+
///
234+
/// Returns a newly allocated set containing the elements in the range
235+
/// `[at, len)`. After the call, the original set will be left containing
236+
/// the elements `[0, at)` with its previous capacity unchanged.
237+
///
238+
/// ***Panics*** if `at > len`.
239+
pub fn split_off(&mut self, at: usize) -> Self
240+
where
241+
S: Clone,
242+
{
243+
Self {
244+
map: self.map.split_off(at),
245+
}
246+
}
224247
}
225248

226249
impl<T, S> IndexSet<T, S>
@@ -576,10 +599,18 @@ impl<T, S> IndexSet<T, S> {
576599
///
577600
/// Computes in **O(1)** time.
578601
pub fn get_index(&self, index: usize) -> Option<&T> {
579-
self.map.get_index(index).map(|(x, &())| x)
602+
self.as_entries().get(index).map(Bucket::key_ref)
603+
}
604+
605+
pub fn first(&self) -> Option<&T> {
606+
self.as_entries().first().map(Bucket::key_ref)
607+
}
608+
609+
pub fn last(&self) -> Option<&T> {
610+
self.as_entries().last().map(Bucket::key_ref)
580611
}
581612

582-
/// Remove the key-value pair by index
613+
/// Remove the value by index
583614
///
584615
/// Valid indices are *0 <= index < self.len()*
585616
///
@@ -592,7 +623,7 @@ impl<T, S> IndexSet<T, S> {
592623
self.map.swap_remove_index(index).map(|(x, ())| x)
593624
}
594625

595-
/// Remove the key-value pair by index
626+
/// Remove the value by index
596627
///
597628
/// Valid indices are *0 <= index < self.len()*
598629
///
@@ -604,6 +635,13 @@ impl<T, S> IndexSet<T, S> {
604635
pub fn shift_remove_index(&mut self, index: usize) -> Option<T> {
605636
self.map.shift_remove_index(index).map(|(x, ())| x)
606637
}
638+
639+
/// Swaps the position of two values in the set.
640+
///
641+
/// ***Panics*** if `a` or `b` are out of bounds.
642+
pub fn swap_indices(&mut self, a: usize, b: usize) {
643+
self.map.swap_indices(a, b)
644+
}
607645
}
608646

609647
/// Access `IndexSet` values at indexed positions.

0 commit comments

Comments
 (0)