Merge pull request #402 from 4e554c4c/insert_by

Add `insert_sorted_by{,_key}` methods for map/set

Author: cuviper

src/map.rs

@@ -493,6 +493,54 @@ where
         }
     }
 
+    /// Insert a key-value pair in the map at its ordered position among keys
+    /// sorted by `cmp`.
+    ///
+    /// This is equivalent to finding the position with
+    /// [`binary_search_by`][Self::binary_search_by], then calling
+    /// [`insert_before`][Self::insert_before] with the given key and 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 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>)
+    where
+        K: Ord,
+        F: FnMut(&K, &V) -> Ordering,
+    {
+        let (Ok(i) | Err(i)) = self.binary_search_by(cmp);
+        self.insert_before(i, key, value)
+    }
+
+    /// Insert a key-value pair in the map at its ordered position
+    /// using a sort-key extraction function.
+    ///
+    /// This is equivalent to finding the position with
+    /// [`binary_search_by_key`][Self::binary_search_by_key] with `sort_key(key)`, then
+    /// calling [`insert_before`][Self::insert_before] with the given key and 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 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,
+        key: K,
+        value: V,
+    ) -> (usize, Option<V>)
+    where
+        B: Ord,
+        F: FnMut(&K, &V) -> B,
+    {
+        let search_key = sort_key(&key, &value);
+        let (Ok(i) | Err(i)) = self.binary_search_by_key(&search_key, sort_key);
+        self.insert_before(i, key, value)
+    }
+
     /// Insert a key-value pair in the map before the entry at the given index, or at the end.
     ///
     /// If an equivalent key already exists in the map: the key remains and

src/map/tests.rs

@@ -1226,3 +1226,41 @@ fn disjoint_indices_mut_fail_duplicate() {
         Err(crate::GetDisjointMutError::OverlappingIndices)
     );
 }
+
+#[test]
+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);
+        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);
+        assert_eq!(old, Some(*value));
+        *value = -*value;
+    }
+    assert_eq!(values, *map.as_slice());
+}
+
+#[test]
+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);
+        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);
+        assert_eq!(old, Some(*value));
+        *value = -*value;
+    }
+    assert_eq!(values, *map.as_slice());
+}

src/set.rs

@@ -422,6 +422,48 @@ where
         (index, existing.is_none())
     }
 
+    /// Insert the value into the set at its ordered position among values
+    /// sorted by `cmp`.
+    ///
+    /// This is equivalent to finding the position with
+    /// [`binary_search_by`][Self::binary_search_by], then calling
+    /// [`insert_before`][Self::insert_before].
+    ///
+    /// If the existing items are **not** already sorted, then the insertion
+    /// index is unspecified (like [`slice::binary_search`]), but the value
+    /// 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)
+    where
+        T: Ord,
+        F: FnMut(&T) -> Ordering,
+    {
+        let (index, existing) = self.map.insert_sorted_by(|k, _| cmp(k), value, ());
+        (index, existing.is_none())
+    }
+
+    /// Insert the value into the set at its ordered position among values
+    /// using a sort-key extraction function.
+    ///
+    /// This is equivalent to finding the position with
+    /// [`binary_search_by_key`][Self::binary_search_by_key] with `sort_key(key)`,
+    /// then calling [`insert_before`][Self::insert_before].
+    ///
+    /// If the existing items are **not** already sorted, then the insertion
+    /// index is unspecified (like [`slice::binary_search`]), but the value
+    /// 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)
+    where
+        B: Ord,
+        F: FnMut(&T) -> B,
+    {
+        let (index, existing) = self.map.insert_sorted_by_key(|k, _| sort_key(k), value, ());
+        (index, existing.is_none())
+    }
+
     /// Insert the value into the set before the value at the given index, or at the end.
     ///
     /// If an equivalent item already exists in the set, it returns `false` leaving the