Merge pull request #406 from cuviper/more-insert_sorted_by

Tweak and extend `insert sorted by`{,`_key`} methods

Author: cuviper

src/map.rs

@@ -505,12 +505,12 @@ where
     /// pair is moved to or inserted at that position regardless.
     ///
     /// Computes in **O(n)** time (average).
-    pub fn insert_sorted_by<F>(&mut self, cmp: F, key: K, value: V) -> (usize, Option<V>)
+    pub fn insert_sorted_by<F>(&mut self, key: K, value: V, mut cmp: F) -> (usize, Option<V>)
     where
         K: Ord,
-        F: FnMut(&K, &V) -> Ordering,
+        F: FnMut(&K, &V, &K, &V) -> Ordering,
     {
-        let (Ok(i) | Err(i)) = self.binary_search_by(cmp);
+        let (Ok(i) | Err(i)) = self.binary_search_by(|k, v| cmp(k, v, &key, &value));
         self.insert_before(i, key, value)
     }
 
@@ -526,11 +526,11 @@ where
     /// pair is moved to or inserted at that position regardless.
     ///
     /// Computes in **O(n)** time (average).
-    pub fn insert_sorted_by_key<F, B>(
+    pub fn insert_sorted_by_key<B, F>(
         &mut self,
-        mut sort_key: F,
         key: K,
         value: V,
+        mut sort_key: F,
     ) -> (usize, Option<V>)
     where
         B: Ord,

src/map/core/entry.rs

@@ -1,5 +1,6 @@
 use super::{equivalent, Entries, IndexMapCore, RefMut};
 use crate::HashValue;
+use core::cmp::Ordering;
 use core::{fmt, mem};
 use hashbrown::hash_table;
 
@@ -402,6 +403,45 @@ impl<'a, K, V> VacantEntry<'a, K, V> {
         (i, self.shift_insert(i, value))
     }
 
+    /// Inserts the entry's key and the given value into the map at its ordered
+    /// position among keys sorted by `cmp`, and returns the new index and a
+    /// mutable reference to the value.
+    ///
+    /// If the existing keys are **not** already sorted, then the insertion
+    /// index is unspecified (like [`slice::binary_search`]), but the key-value
+    /// pair is inserted at that position regardless.
+    ///
+    /// Computes in **O(n)** time (average).
+    pub fn insert_sorted_by<F>(self, value: V, mut cmp: F) -> (usize, &'a mut V)
+    where
+        K: Ord,
+        F: FnMut(&K, &V, &K, &V) -> Ordering,
+    {
+        let slice = crate::map::Slice::from_slice(self.map.entries);
+        let (Ok(i) | Err(i)) = slice.binary_search_by(|k, v| cmp(k, v, &self.key, &value));
+        (i, self.shift_insert(i, value))
+    }
+
+    /// Inserts the entry's key and the given value into the map at its ordered
+    /// position using a sort-key extraction function, and returns the new index
+    /// and a mutable reference to the value.
+    ///
+    /// If the existing keys are **not** already sorted, then the insertion
+    /// index is unspecified (like [`slice::binary_search`]), but the key-value
+    /// pair is inserted at that position regardless.
+    ///
+    /// Computes in **O(n)** time (average).
+    pub fn insert_sorted_by_key<B, F>(self, value: V, mut sort_key: F) -> (usize, &'a mut V)
+    where
+        B: Ord,
+        F: FnMut(&K, &V) -> B,
+    {
+        let search_key = sort_key(&self.key, &value);
+        let slice = crate::map::Slice::from_slice(self.map.entries);
+        let (Ok(i) | Err(i)) = slice.binary_search_by_key(&search_key, sort_key);
+        (i, self.shift_insert(i, value))
+    }
+
     /// Inserts the entry's key and the given value into the map at the given index,
     /// shifting others to the right, and returns a mutable reference to the value.
     ///

src/map/tests.rs

@@ -1232,14 +1232,14 @@ fn insert_sorted_by_key() {
     let mut values = [(-1, 8), (3, 18), (-27, 2), (-2, 5)];
     let mut map: IndexMap<i32, i32> = IndexMap::new();
     for (key, value) in values {
-        let (_, old) = map.insert_sorted_by_key(|k, _| k.abs(), key, value);
+        let (_, old) = map.insert_sorted_by_key(key, value, |k, _| k.abs());
         assert_eq!(old, None);
     }
     values.sort_by_key(|(key, _)| key.abs());
     assert_eq!(values, *map.as_slice());
 
     for (key, value) in &mut values {
-        let (_, old) = map.insert_sorted_by_key(|k, _| k.abs(), *key, -*value);
+        let (_, old) = map.insert_sorted_by_key(*key, -*value, |k, _| k.abs());
         assert_eq!(old, Some(*value));
         *value = -*value;
     }
@@ -1251,14 +1251,14 @@ fn insert_sorted_by() {
     let mut values = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)];
     let mut map: IndexMap<i32, i32> = IndexMap::new();
     for (key, value) in values {
-        let (_, old) = map.insert_sorted_by(|probe, _| key.cmp(probe), key, value);
+        let (_, old) = map.insert_sorted_by(key, value, |key1, _, key2, _| key2.cmp(key1));
         assert_eq!(old, None);
     }
     values.reverse();
     assert_eq!(values, *map.as_slice());
 
     for (key, value) in &mut values {
-        let (_, old) = map.insert_sorted_by(|probe, _| (*key).cmp(probe), *key, -*value);
+        let (_, old) = map.insert_sorted_by(*key, -*value, |key1, _, key2, _| key2.cmp(key1));
         assert_eq!(old, Some(*value));
         *value = -*value;
     }

src/set.rs

@@ -434,12 +434,14 @@ where
     /// is moved to or inserted at that position regardless.
     ///
     /// Computes in **O(n)** time (average).
-    pub fn insert_sorted_by<F>(&mut self, mut cmp: F, value: T) -> (usize, bool)
+    pub fn insert_sorted_by<F>(&mut self, value: T, mut cmp: F) -> (usize, bool)
     where
         T: Ord,
-        F: FnMut(&T) -> Ordering,
+        F: FnMut(&T, &T) -> Ordering,
     {
-        let (index, existing) = self.map.insert_sorted_by(|k, _| cmp(k), value, ());
+        let (index, existing) = self
+            .map
+            .insert_sorted_by(value, (), |a, (), b, ()| cmp(a, b));
         (index, existing.is_none())
     }
 
@@ -455,12 +457,12 @@ where
     /// is moved to or inserted at that position regardless.
     ///
     /// Computes in **O(n)** time (average).
-    pub fn insert_sorted_by_key<F, B>(&mut self, mut sort_key: F, value: T) -> (usize, bool)
+    pub fn insert_sorted_by_key<B, F>(&mut self, value: T, mut sort_key: F) -> (usize, bool)
     where
         B: Ord,
         F: FnMut(&T) -> B,
     {
-        let (index, existing) = self.map.insert_sorted_by_key(|k, _| sort_key(k), value, ());
+        let (index, existing) = self.map.insert_sorted_by_key(value, (), |k, _| sort_key(k));
         (index, existing.is_none())
     }
 
@@ -938,7 +940,7 @@ impl<T, S> IndexSet<T, S> {
     where
         F: FnMut(&T, &T) -> Ordering,
     {
-        self.map.sort_by(move |a, _, b, _| cmp(a, b));
+        self.map.sort_by(move |a, (), b, ()| cmp(a, b));
     }
 
     /// Sort the values of the set and return a by-value iterator of

tests/quick.rs

@@ -130,6 +130,46 @@ quickcheck_limit! {
         true
     }
 
+    fn insert_sorted_by(insert: Vec<(u32, u32)>) -> bool {
+        let mut hmap = HashMap::new();
+        let mut map = IndexMap::new();
+        let mut map2 = IndexMap::new();
+        for &(key, value) in &insert {
+            hmap.insert(key, value);
+            map.insert_sorted_by(key, value, |key1, _, key2, _| key2.cmp(key1));
+            match map2.entry(key) {
+                Entry::Occupied(e) => *e.into_mut() = value,
+                Entry::Vacant(e) => {
+                    e.insert_sorted_by(value, |key1, _, key2, _| key2.cmp(key1));
+                }
+            }
+        }
+        let hsorted = hmap.iter().sorted_by(|(key1, _), (key2, _)| key2.cmp(key1));
+        itertools::assert_equal(hsorted, &map);
+        itertools::assert_equal(&map, &map2);
+        true
+    }
+
+    fn insert_sorted_by_key(insert: Vec<(i32, u32)>) -> bool {
+        let mut hmap = HashMap::new();
+        let mut map = IndexMap::new();
+        let mut map2 = IndexMap::new();
+        for &(key, value) in &insert {
+            hmap.insert(key, value);
+            map.insert_sorted_by_key(key, value, |&k, _| (k.unsigned_abs(), k));
+            match map2.entry(key) {
+                Entry::Occupied(e) => *e.into_mut() = value,
+                Entry::Vacant(e) => {
+                    e.insert_sorted_by_key(value, |&k, _| (k.unsigned_abs(), k));
+                }
+            }
+        }
+        let hsorted = hmap.iter().sorted_by_key(|(&k, _)| (k.unsigned_abs(), k));
+        itertools::assert_equal(hsorted, &map);
+        itertools::assert_equal(&map, &map2);
+        true
+    }
+
     fn replace_index(insert: Vec<u8>, index: u8, new_key: u8) -> TestResult {
         if insert.is_empty() {
             return TestResult::discard();