From 6b54affc24043191eec984ac27a53dcaa3ab0598 Mon Sep 17 00:00:00 2001 From: William Rieger Date: Thu, 11 Feb 2021 14:20:04 -0500 Subject: [PATCH 1/4] Add `BinaryHeap::contains` and `BinaryHeap::remove` --- library/alloc/src/collections/binary_heap.rs | 89 ++++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/library/alloc/src/collections/binary_heap.rs b/library/alloc/src/collections/binary_heap.rs index 8a36b2af76522..8267fbe95e8db 100644 --- a/library/alloc/src/collections/binary_heap.rs +++ b/library/alloc/src/collections/binary_heap.rs @@ -143,6 +143,8 @@ #![allow(missing_docs)] #![stable(feature = "rust1", since = "1.0.0")] +use core::borrow::Borrow; +use core::cmp::Ordering; use core::fmt; use core::iter::{FromIterator, FusedIterator, InPlaceIterable, SourceIter, TrustedLen}; use core::mem::{self, swap, ManuallyDrop}; @@ -708,6 +710,93 @@ impl BinaryHeap { self.data.retain(f); self.rebuild(); } + + /// Returns `true` if the `BinaryHeap` contains the specified item. + /// + /// The item may be any borrowed form of the heap's item type, but + /// `PartialEq` on the borrowed form *must* match that for the item type. + /// + /// # Examples + /// + /// Basic usage: + /// + /// ``` + /// #![feature(binary_heap_contains_remove)] + /// use std::collections::BinaryHeap; + /// + /// let heap = BinaryHeap::from(vec![1, 3, 5, 7, 9]); + /// + /// assert!(heap.contains(&5)); + /// assert!(!heap.contains(&2)); + /// ``` + /// + /// # Time complexity + /// + /// The worst case time complexity is *O*(*n*). + #[unstable(feature = "binary_heap_contains_remove", issue = "82001")] + pub fn contains(&self, item: &Q) -> bool + where + T: Borrow, + Q: PartialEq, + { + self.data.iter().any(|x| x.borrow() == item) + } + + /// Removes an item from the `BinaryHeap` and returns it, if it exists. + /// + /// The item may be any borrowed form of the heap's item type, but + /// `PartialEq` on the borrowed form *must* match that for the item type. If + /// there are multiple matching elements, an arbitrary and unspecified one is + /// removed. + /// + /// # Examples + /// + /// Basic usage: + /// + /// ``` + /// #![feature(binary_heap_contains_remove)] + /// use std::collections::BinaryHeap; + /// + /// let mut heap = BinaryHeap::from(vec![1, 3, 5, 7, 9]); + /// + /// assert_eq!(heap.remove(&7), Some(7)); + /// assert_eq!(heap.remove(&7), None); + /// ``` + /// + /// # Time complexity + /// + /// The worst case time complexity is *O*(*n*). + #[unstable(feature = "binary_heap_contains_remove", issue = "82001")] + pub fn remove(&mut self, item: &Q) -> Option + where + T: Borrow, + Q: PartialEq, + { + self.data.iter().position(|x| x.borrow() == item).map(|pos| { + let last = self.data.len() - 1; + let comparison = self.data[last].cmp(&self.data[pos]); + let removed_item = self.data.swap_remove(pos); + match comparison { + Ordering::Less => { + // If the swapped item is less than the removed item, sift + // it down the heap to its correct position. + debug_assert!(pos < last); + self.sift_down(pos); + } + Ordering::Greater => { + // If the swapped item is greater than the removed item, + // sift it up the heap to its correct position. + debug_assert!(pos < last); + self.sift_up(0, pos); + } + Ordering::Equal => { + // If the swapped item is equal to the removed item, it's + // already in the correct position. + } + } + removed_item + }) + } } impl BinaryHeap { From aa4497d6cdf0a64f0a8e03e2755da670f7284296 Mon Sep 17 00:00:00 2001 From: William Rieger Date: Thu, 11 Feb 2021 14:20:19 -0500 Subject: [PATCH 2/4] Add tests --- library/alloc/tests/binary_heap.rs | 36 ++++++++++++++++++++++++++++++ library/alloc/tests/lib.rs | 1 + 2 files changed, 37 insertions(+) diff --git a/library/alloc/tests/binary_heap.rs b/library/alloc/tests/binary_heap.rs index ce794a9a4afa2..907b639fbe90f 100644 --- a/library/alloc/tests/binary_heap.rs +++ b/library/alloc/tests/binary_heap.rs @@ -392,6 +392,42 @@ fn test_retain() { assert_eq!(a.into_sorted_vec(), [-10, 2, 4]) } +#[test] +fn test_contains_remove() { + let mut a = BinaryHeap::from(vec![5, 7, 13, -2, -2, 8, -9]); + + assert!(a.contains(&7)); + assert_eq!(a.remove(&7), Some(7)); + assert_eq!(a.remove(&7), None); + assert!(!a.contains(&7)); + + assert!(a.contains(&-2)); + assert_eq!(a.remove(&-2), Some(-2)); + assert_eq!(a.remove(&-2), Some(-2)); + assert_eq!(a.remove(&-2), None); + assert!(!a.contains(&-2)); + + assert!(a.contains(&13)); + assert_eq!(a.remove(&13), Some(13)); + assert_eq!(a.remove(&13), None); + assert!(!a.contains(&13)); + + assert_eq!(a.into_sorted_vec(), [-9, 5, 8]); +} + +#[test] +fn test_empty_remove() { + let mut a = BinaryHeap::::new(); + assert_eq!(a.remove(&5), None); +} + +#[test] +fn test_singleton_remove() { + let mut a = BinaryHeap::from(vec![5]); + assert_eq!(a.remove(&5), Some(5)); + assert!(a.is_empty()); +} + // old binaryheap failed this test // // Integrity means that all elements are present after a comparison panics, diff --git a/library/alloc/tests/lib.rs b/library/alloc/tests/lib.rs index dd98f806451d8..26167f6fb6afd 100644 --- a/library/alloc/tests/lib.rs +++ b/library/alloc/tests/lib.rs @@ -21,6 +21,7 @@ #![feature(vecdeque_binary_search)] #![feature(slice_group_by)] #![feature(vec_extend_from_within)] +#![feature(binary_heap_contains_remove)] use std::collections::hash_map::DefaultHasher; use std::hash::{Hash, Hasher}; From 2acc337c055bb98350dcc379146e94ac1659d96a Mon Sep 17 00:00:00 2001 From: William Rieger Date: Thu, 11 Feb 2021 22:09:42 -0500 Subject: [PATCH 3/4] Doc comment improvements --- library/alloc/src/collections/binary_heap.rs | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/library/alloc/src/collections/binary_heap.rs b/library/alloc/src/collections/binary_heap.rs index 8267fbe95e8db..cf144940ee96d 100644 --- a/library/alloc/src/collections/binary_heap.rs +++ b/library/alloc/src/collections/binary_heap.rs @@ -724,10 +724,13 @@ impl BinaryHeap { /// #![feature(binary_heap_contains_remove)] /// use std::collections::BinaryHeap; /// - /// let heap = BinaryHeap::from(vec![1, 3, 5, 7, 9]); + /// let mut heap = BinaryHeap::new(); + /// heap.push("dog".to_string()); + /// heap.push("cat".to_string()); + /// heap.push("fish".to_string()); /// - /// assert!(heap.contains(&5)); - /// assert!(!heap.contains(&2)); + /// assert!(heap.contains("cat")); + /// assert!(!heap.contains("horse")); /// ``` /// /// # Time complexity @@ -746,8 +749,8 @@ impl BinaryHeap { /// /// The item may be any borrowed form of the heap's item type, but /// `PartialEq` on the borrowed form *must* match that for the item type. If - /// there are multiple matching elements, an arbitrary and unspecified one is - /// removed. + /// there are multiple matching elements, there are no guarantees which of + /// the matching elements is removed on each call to `remove()`. /// /// # Examples /// From 763011766aea5e32a268095cea1d7699ff87e1a3 Mon Sep 17 00:00:00 2001 From: William Rieger Date: Fri, 12 Feb 2021 00:26:11 -0500 Subject: [PATCH 4/4] Change debug assertions to actual asserions I think incorrect (or malicious) `PartialEq` and `PartialOrd` implementations for `T` could define `x == y` and `x < y` to both be conditionally true with interior mutability in such a way that `pos == len - 1` (implying `data[pos] == data[len - 1]`) but we end up in the `data[pos] < data[len - 1]` branch. `sift_up()` and `sift_down()` assume that `pos` is in-bounds which would be violated in this case. --- library/alloc/src/collections/binary_heap.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/library/alloc/src/collections/binary_heap.rs b/library/alloc/src/collections/binary_heap.rs index cf144940ee96d..41df37ebd8d3a 100644 --- a/library/alloc/src/collections/binary_heap.rs +++ b/library/alloc/src/collections/binary_heap.rs @@ -783,13 +783,13 @@ impl BinaryHeap { Ordering::Less => { // If the swapped item is less than the removed item, sift // it down the heap to its correct position. - debug_assert!(pos < last); + assert!(pos < last); self.sift_down(pos); } Ordering::Greater => { // If the swapped item is greater than the removed item, // sift it up the heap to its correct position. - debug_assert!(pos < last); + assert!(pos < last); self.sift_up(0, pos); } Ordering::Equal => {