Add more Vec/slice-like methods to maps and sets

src/map.rs

@@ -252,6 +252,13 @@ impl<K, V, S> IndexMap<K, V, S> {
         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<K, V, S> IndexMap<K, V, S> {
             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<K, V, S> IndexMap<K, V, S>
@@ -693,6 +717,34 @@ impl<K, V, S> IndexMap<K, V, S> {
         self.as_entries_mut().get_mut(index).map(Bucket::muts)
     }
 
+    /// Get the first key-value pair
+    ///
+    /// Computes in **O(1)** time.
+    pub fn first(&self) -> Option<(&K, &V)> {
+        self.as_entries().first().map(Bucket::refs)
+    }
+
+    /// Get the first key-value pair, with mutable access to the value
+    ///
+    /// Computes in **O(1)** time.
+    pub fn first_mut(&mut self) -> Option<(&K, &mut V)> {
+        self.as_entries_mut().first_mut().map(Bucket::ref_mut)
+    }
+
+    /// Get the last key-value pair
+    ///
+    /// Computes in **O(1)** time.
+    pub fn last(&self) -> Option<(&K, &V)> {
+        self.as_entries().last().map(Bucket::refs)
+    }
+
+    /// Get the last key-value pair, with mutable access to the value
+    ///
+    /// Computes in **O(1)** time.
+    pub fn last_mut(&mut self) -> Option<(&K, &mut V)> {
+        self.as_entries_mut().last_mut().map(Bucket::ref_mut)
+    }
+
     /// Remove the key-value pair by index
     ///
     /// Valid indices are *0 <= index < self.len()*
@@ -718,6 +770,13 @@ impl<K, V, S> IndexMap<K, V, S> {
     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`.

src/map/core.rs

@@ -150,6 +150,13 @@ impl<K, V> IndexMapCore<K, V> {
         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<R>(&mut self, range: R) -> Drain<'_, Bucket<K, V>>
     where
         R: RangeBounds<usize>,
@@ -159,6 +166,18 @@ impl<K, V> IndexMapCore<K, V> {
         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));

src/map/core/raw.rs

@@ -61,6 +61,29 @@ impl<K, V> IndexMapCore<K, V> {
         // 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`.

src/set.rs

@@ -200,6 +200,13 @@ impl<T, S> IndexSet<T, S> {
         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<T, S> IndexSet<T, S> {
             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<T, S> IndexSet<T, S>
@@ -576,10 +599,24 @@ impl<T, S> IndexSet<T, S> {
     ///
     /// 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)
     }
 
-    /// Remove the key-value pair by index
+    /// Get the first value
+    ///
+    /// Computes in **O(1)** time.
+    pub fn first(&self) -> Option<&T> {
+        self.as_entries().first().map(Bucket::key_ref)
+    }
+
+    /// Get the last value
+    ///
+    /// Computes in **O(1)** time.
+    pub fn last(&self) -> Option<&T> {
+        self.as_entries().last().map(Bucket::key_ref)
+    }
+
+    /// Remove the value by index
     ///
     /// Valid indices are *0 <= index < self.len()*
     ///
@@ -592,7 +629,7 @@ impl<T, S> IndexSet<T, S> {
         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 +641,13 @@ impl<T, S> IndexSet<T, S> {
     pub fn shift_remove_index(&mut self, index: usize) -> Option<T> {
         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.