Add Heap.removeAll(where:)

Sources/HeapModule/Heap.swift

@@ -268,6 +268,28 @@ extension Heap {
     _checkInvariants()
     return removed
   }
+
+  /// Removes all the elements that satisfy the given predicate.
+  ///
+  /// - 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 {
+    defer {
+      if _storage.count > 1 {
+        _update { handle in
+          handle.heapify()
+        }
+      }
+      _checkInvariants()
+    }
+    try _storage.removeAll(where: shouldBeRemoved)
+  }
 
   /// Replaces the maximum value in the heap with the given replacement,
   /// then updates heap contents to reflect the change.

Tests/HeapTests/HeapTests.swift

@@ -641,4 +641,67 @@ final class HeapTests: CollectionTestCase {
       }
     }
   }
+
+  func test_removeAll_noneRemoved() {
+    withEvery("count", in: 0 ..< 20) { count in
+      withEvery("seed", in: 0 ..< 10) { seed in
+        var rng = RepeatableRandomNumberGenerator(seed: seed)
+        let input = (0 ..< count).shuffled(using: &rng)
+        var heap = Heap(input)
+        heap.removeAll { _ in false }
+        let expected = Array(0 ..< count)
+        expectEqualElements(heap.itemsInAscendingOrder(), expected)
+      }
+    }
+  }
+
+  func test_removeAll_allRemoved() {
+    withEvery("count", in: 0 ..< 20) { count in
+      withEvery("seed", in: 0 ..< 10) { seed in
+        var rng = RepeatableRandomNumberGenerator(seed: seed)
+        let input = (0 ..< count).shuffled(using: &rng)
+        var heap = Heap(input)
+        heap.removeAll { _ in true }
+        expectTrue(heap.isEmpty)
+      }
+    }
+  }
+
+  func test_removeAll_removeEvenNumbers() {
+    withEvery("count", in: 0 ..< 20) { count in
+      withEvery("seed", in: 0 ..< 10) { seed in
+        var rng = RepeatableRandomNumberGenerator(seed: seed)
+        let input = (0 ..< count).shuffled(using: &rng)
+        var heap = Heap(input)
+        heap.removeAll { $0 % 2 == 0 }
+        let expected = Array(stride(from: 1, to: count, by: 2))
+        expectEqualElements(heap.itemsInAscendingOrder(), expected)
+      }
+    }
+  }
+
+  func test_removeAll_throw() throws {
+    struct DummyError: Error {}
+
+    try withEvery("count", in: 1 ..< 20) { count in
+      try withEvery("seed", in: 0 ..< 10) { seed in
+        var rng = RepeatableRandomNumberGenerator(seed: seed)
+        let input = (0 ..< count).shuffled(using: &rng)
+        var heap = Heap(input)
+        expectThrows(
+          try heap.removeAll { v in
+            if v == count / 2 {
+              throw DummyError()
+            }
+            return v % 2 == 0
+          }
+        ) { error in
+          expectTrue(error is DummyError)
+        }
+        // Throwing halfway through `removeAll` is expected to reorder items,
+        // but not remove any.
+        expectEqualElements(heap.itemsInAscendingOrder(), 0 ..< count)
+      }
+    }
+  }
 }