Add Heap.removeAll(where:)

Author: DakshinD

Sources/HeapModule/Heap.swift

@@ -268,6 +268,30 @@ extension Heap {
     _checkInvariants()
     return removed
   }
+    
+  /// Removes all the elements that satisfy the given predicate.
+  ///
+  /// The heap *must not* be empty.
+  ///
+  /// - Parameter shouldBeRemoved: A closure that takes an element of the
+  ///   heap as its argument and returns a Boolean value indicating
+  ///   whether the element should be removed from the heap.
+  ///
+  /// - Complexity: O(*n*), where *n* is the number of items in the heap.
+  @inlinable
+  public mutating func removeAll(
+      where shouldBeRemoved: (Element) throws -> Bool
+  ) rethrows {
+      // Precondition: the heap must not be empty.
+      precondition(!isEmpty, "Heap must not be empty")
+      _storage = try _storage.filter { try !shouldBeRemoved($0) }
+      if _storage.count > 1 {
+      _update { handle in
+          handle.heapify()
+      }
+      }
+      _checkInvariants()
+  }
 
   /// Replaces the maximum value in the heap with the given replacement,
   /// then updates heap contents to reflect the change.

Tests/HeapTests/HeapTests.swift

@@ -587,4 +587,24 @@ final class HeapTests: CollectionTestCase {
       }
     }
   }
+    
+  func test_removeAll_noneRemoved() {
+      var heap = Heap<Int>((1...10).shuffled())
+      let originalSorted = heap.unordered.sorted()
+      heap.removeAll { _ in false }
+      expectEqualElements(heap.itemsInAscendingOrder(), originalSorted)
+  }
+    
+  func test_removeAll_allRemoved() {
+      var heap = Heap<Int>((1...10).shuffled())
+      heap.removeAll { _ in true }
+      expectTrue(heap.isEmpty)
+  }
+    
+  func test_removeAll_removeEvenNumbers() {
+      var heap = Heap<Int>((1...10).shuffled())
+      heap.removeAll { $0 % 2 == 0 }
+      let expected = (1...10).filter { $0 % 2 != 0 }.sorted()
+      expectEqualElements(heap.itemsInAscendingOrder(), expected)
+  }
 }