From 3227a23b008c602c171f44a764da7ddc0d448527 Mon Sep 17 00:00:00 2001 From: Josh Stone Date: Sat, 24 Feb 2024 15:56:15 -0800 Subject: [PATCH] Add `insert_sorted` --- src/map.rs | 33 +++++++++++++++++++++++++++++++++ src/map/core/entry.rs | 18 ++++++++++++++++++ src/set.rs | 30 ++++++++++++++++++++++++++++++ tests/quick.rs | 17 +++++++++++++++++ 4 files changed, 98 insertions(+) diff --git a/src/map.rs b/src/map.rs index 676b31aa..a060fad1 100644 --- a/src/map.rs +++ b/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) + 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 @@ -774,6 +803,10 @@ impl IndexMap { /// Sort the map’s key-value pairs by the default ordering of the keys. /// + /// This is a stable sort -- but since equivalent keys should not normally coexist in + /// a map at all, [`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 diff --git a/src/map/core/entry.rs b/src/map/core/entry.rs index 48c00294..6c310707 100644 --- a/src/map/core/entry.rs +++ b/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. /// diff --git a/src/set.rs b/src/set.rs index 0c0358cd..9f10e3ba 100644 --- a/src/set.rs +++ b/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 IndexSet { /// Sort the set’s values by their default ordering. /// + /// This is a stable sort -- but since equivalent values should not normally coexist + /// in a set at all, [`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 diff --git a/tests/quick.rs b/tests/quick.rs index c32d0b91..56afee72 100644 --- a/tests/quick.rs +++ b/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) -> bool { let mut map = IndexMap::new(); for &key in &insert {