Add sort_by_cached_key and par_sort_by_cached_key

Author: cuviper

Cargo.toml

@@ -18,7 +18,7 @@ bench = false
 arbitrary = { version = "1.0", optional = true, default-features = false }
 quickcheck = { version = "1.0", optional = true, default-features = false }
 serde = { version = "1.0", optional = true, default-features = false }
-rayon = { version = "1.4.1", optional = true }
+rayon = { version = "1.5.3", optional = true }
 
 # Internal feature, only used when building as part of rustc,
 # not part of the stable interface of this crate.

RELEASES.md

@@ -20,6 +20,10 @@
   - `IndexMap` and `IndexSet` both implement `arbitrary::Arbitrary<'_>` and
     `quickcheck::Arbitrary` if those optional dependency features are enabled.
 
+  - `IndexMap` and `IndexSet` now have `sort_by_cached_key` and
+    `par_sort_by_cached_key` methods which perform stable sorts in place
+    using a key extraction function.
+
   - The `hashbrown` dependency has been updated to version 0.13.
 
 - 1.9.1

src/map.rs

@@ -736,6 +736,24 @@ where
         IntoIter::new(entries)
     }
 
+    /// Sort the map’s key-value pairs in place using a sort-key extraction function.
+    ///
+    /// During sorting, the function is called at most once per entry, by using temporary storage
+    /// to remember the results of its evaluation. The order of calls to the function is
+    /// unspecified and may change between versions of `indexmap` or the standard library.
+    ///
+    /// Computes in **O(m n + n log n + c)** time () and **O(n)** space, where the function is
+    /// **O(m)**, *n* is the length of the map, and *c* the capacity. The sort is stable.
+    pub fn sort_by_cached_key<T, F>(&mut self, mut sort_key: F)
+    where
+        T: Ord,
+        F: FnMut(&K, &V) -> T,
+    {
+        self.with_entries(move |entries| {
+            entries.sort_by_cached_key(move |a| sort_key(&a.key, &a.value));
+        });
+    }
+
     /// Reverses the order of the map’s key-value pairs in place.
     ///
     /// Computes in **O(n)** time and **O(1)** space.

src/rayon/map.rs

@@ -504,6 +504,18 @@ where
         entries.par_sort_unstable_by(move |a, b| cmp(&a.key, &a.value, &b.key, &b.value));
         IntoParIter { entries }
     }
+
+    /// Sort the map’s key-value pairs in place and in parallel, using a sort-key extraction
+    /// function.
+    pub fn par_sort_by_cached_key<T, F>(&mut self, sort_key: F)
+    where
+        T: Ord + Send,
+        F: Fn(&K, &V) -> T + Sync,
+    {
+        self.with_entries(move |entries| {
+            entries.par_sort_by_cached_key(move |a| sort_key(&a.key, &a.value));
+        });
+    }
 }
 
 /// A parallel mutable iterator over the values of a `IndexMap`.

src/rayon/set.rs

@@ -572,6 +572,17 @@ where
         entries.par_sort_unstable_by(move |a, b| cmp(&a.key, &b.key));
         IntoParIter { entries }
     }
+
+    /// Sort the set’s values in place and in parallel, using a key extraction function.
+    pub fn par_sort_by_cached_key<K, F>(&mut self, sort_key: F)
+    where
+        K: Ord + Send,
+        F: Fn(&T) -> K + Sync,
+    {
+        self.with_entries(move |entries| {
+            entries.par_sort_by_cached_key(move |a| sort_key(&a.key));
+        });
+    }
 }
 
 /// Requires crate feature `"rayon"`.

src/set.rs

@@ -633,6 +633,24 @@ where
         IntoIter::new(entries)
     }
 
+    /// Sort the set’s values in place using a key extraction function.
+    ///
+    /// During sorting, the function is called at most once per entry, by using temporary storage
+    /// to remember the results of its evaluation. The order of calls to the function is
+    /// unspecified and may change between versions of `indexmap` or the standard library.
+    ///
+    /// Computes in **O(m n + n log n + c)** time () and **O(n)** space, where the function is
+    /// **O(m)**, *n* is the length of the map, and *c* the capacity. The sort is stable.
+    pub fn sort_by_cached_key<K, F>(&mut self, mut sort_key: F)
+    where
+        K: Ord,
+        F: FnMut(&T) -> K,
+    {
+        self.with_entries(move |entries| {
+            entries.sort_by_cached_key(move |a| sort_key(&a.key));
+        });
+    }
+
     /// Reverses the order of the set’s values in place.
     ///
     /// Computes in **O(n)** time and **O(1)** space.

tests/quick.rs

@@ -452,6 +452,12 @@ quickcheck_limit! {
         assert_sorted_by_key(map, |t| t.1);
     }
 
+    fn sort_3(keyvals: Large<Vec<(i8, i8)>>) -> () {
+        let mut map: IndexMap<_, _> = IndexMap::from_iter(keyvals.to_vec());
+        map.sort_by_cached_key(|&k, _| std::cmp::Reverse(k));
+        assert_sorted_by_key(map, |t| std::cmp::Reverse(t.0));
+    }
+
     fn reverse(keyvals: Large<Vec<(i8, i8)>>) -> () {
         let mut map: IndexMap<_, _> = IndexMap::from_iter(keyvals.to_vec());