Merge pull request #315 from cuviper/insert_sorted

Add `insert_sorted`

src/map.rs

@@ -414,6 +414,35 @@ where
         self.core.insert_full(hash, key, value)
     }
 
+    /// Insert a key-value pair in the map at its ordered position among sorted keys.
+    ///
+    /// This is equivalent to finding the position with
+    /// [`binary_search_keys`][Self::binary_search_keys], then either updating
+    /// it or calling [`shift_insert`][Self::shift_insert] for a new key.
+    ///
+    /// If the sorted key is found in the map, its corresponding value is
+    /// updated with `value`, and the older value is returned inside
+    /// `(index, Some(_))`. Otherwise, the new key-value pair is inserted at
+    /// the sorted position, and `(index, None)` is returned.
+    ///
+    /// 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). Instead of repeating calls to
+    /// `insert_sorted`, it may be faster to call batched [`insert`][Self::insert]
+    /// or [`extend`][Self::extend] and only call [`sort_keys`][Self::sort_keys]
+    /// or [`sort_unstable_keys`][Self::sort_unstable_keys] once.
+    pub fn insert_sorted(&mut self, key: K, value: V) -> (usize, Option<V>)
+    where
+        K: Ord,
+    {
+        match self.binary_search_keys(&key) {
+            Ok(i) => (i, Some(mem::replace(&mut self[i], value))),
+            Err(i) => (i, self.shift_insert(i, key, value)),
+        }
+    }
+
     /// Insert a key-value pair in the map at the given index.
     ///
     /// If an equivalent key already exists in the map: the key remains and
@@ -788,6 +817,10 @@ impl<K, V, S> IndexMap<K, V, S> {
 
     /// Sort the map’s key-value pairs by the default ordering of the keys.
     ///
+    /// This is a stable sort -- but equivalent keys should not normally coexist in
+    /// a map at all, so [`sort_unstable_keys`][Self::sort_unstable_keys] is preferred
+    /// because it is generally faster and doesn't allocate auxiliary memory.
+    ///
     /// See [`sort_by`](Self::sort_by) for details.
     pub fn sort_keys(&mut self)
     where

src/map/core/entry.rs

@@ -310,6 +310,24 @@ impl<'a, K, V> VacantEntry<'a, K, V> {
         &mut map.entries[i].value
     }
 
+    /// Inserts the entry's key and the given value into the map at its ordered
+    /// position among sorted keys, 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(self, value: V) -> (usize, &'a mut V)
+    where
+        K: Ord,
+    {
+        let slice = crate::map::Slice::from_slice(&self.map.entries);
+        let i = slice.binary_search_keys(&self.key).unwrap_err();
+        (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/set.rs

@@ -355,6 +355,32 @@ where
         (index, existing.is_none())
     }
 
+    /// Insert the value into the set at its ordered position among sorted values.
+    ///
+    /// This is equivalent to finding the position with
+    /// [`binary_search`][Self::binary_search], and if needed calling
+    /// [`shift_insert`][Self::shift_insert] for a new value.
+    ///
+    /// If the sorted item is found in the set, it returns the index of that
+    /// existing item and `false`, without any change. Otherwise, it inserts the
+    /// new item and returns its sorted index and `true`.
+    ///
+    /// 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). Instead of repeating calls to
+    /// `insert_sorted`, it may be faster to call batched [`insert`][Self::insert]
+    /// or [`extend`][Self::extend] and only call [`sort`][Self::sort] or
+    /// [`sort_unstable`][Self::sort_unstable] once.
+    pub fn insert_sorted(&mut self, value: T) -> (usize, bool)
+    where
+        T: Ord,
+    {
+        let (index, existing) = self.map.insert_sorted(value, ());
+        (index, existing.is_none())
+    }
+
     /// Insert the value into the set at the given index.
     ///
     /// If an equivalent item already exists in the set, it returns
@@ -670,6 +696,10 @@ impl<T, S> IndexSet<T, S> {
 
     /// Sort the set’s values by their default ordering.
     ///
+    /// This is a stable sort -- but equivalent values should not normally coexist in
+    /// a set at all, so [`sort_unstable`][Self::sort_unstable] is preferred
+    /// because it is generally faster and doesn't allocate auxiliary memory.
+    ///
     /// See [`sort_by`](Self::sort_by) for details.
     pub fn sort(&mut self)
     where

tests/quick.rs

@@ -113,6 +113,23 @@ quickcheck_limit! {
         true
     }
 
+    fn insert_sorted(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(key, value);
+            match map2.entry(key) {
+                Entry::Occupied(e) => *e.into_mut() = value,
+                Entry::Vacant(e) => { e.insert_sorted(value); }
+            }
+        }
+        itertools::assert_equal(hmap.iter().sorted(), &map);
+        itertools::assert_equal(&map, &map2);
+        true
+    }
+
     fn pop(insert: Vec<u8>) -> bool {
         let mut map = IndexMap::new();
         for &key in &insert {