[Heap] Rework removals

Author: lorentey

Sources/PriorityQueueModule/Heap+UnsafeHandle.swift

@@ -63,6 +63,13 @@ extension Heap._UnsafeHandle {
   internal func swapAt(_ i: _Node, _ j: _Node) {
     buffer.swapAt(i.offset, j.offset)
   }
+
+  /// Swaps the element at the given node with the supplied value.
+  @inlinable @inline(__always)
+  internal func swapAt(_ i: _Node, with value: inout Element) {
+    let p = buffer.baseAddress.unsafelyUnwrapped + i.offset
+    swap(&p.pointee, &value)
+  }
 }
 
 extension Heap._UnsafeHandle {
@@ -132,19 +139,10 @@ extension Heap._UnsafeHandle {
 }
 
 extension Heap._UnsafeHandle {
-  @inline(__always)
-  @inlinable
-  internal func trickleDown(_ node: _Node) {
-    if node.isMinLevel {
-      trickleDownMin(node)
-    } else {
-      trickleDownMax(node)
-    }
-  }
-
-  @inline(__always)
+  @_effects(releasenone)
   @inlinable
   internal func trickleDownMin(_ node: _Node) {
+    assert(node.isMinLevel)
     var node = node
 
     while let minDescendant = minChildOrGrandchild(of: node) {
@@ -167,9 +165,10 @@ extension Heap._UnsafeHandle {
     }
   }
 
-  @inline(__always)
+  @_effects(releasenone)
   @inlinable
   internal func trickleDownMax(_ node: _Node) {
+    assert(!node.isMinLevel)
     var node = node
 
     while let maxDescendant = maxChildOrGrandchild(of: node) {

Sources/PriorityQueueModule/Heap.swift

@@ -106,28 +106,47 @@ public struct Heap<Element: Comparable> {
   /// - Complexity: O(log `count`)
   @inlinable
   public mutating func popMin() -> Element? {
-    defer { _checkInvariants() }
-    return _remove(.root)
+    guard _storage.count > 0 else { return nil }
+
+    var removed = _storage.removeLast()
+
+    if _storage.count > 0 {
+      _update { handle in
+        let minNode = _Node.root
+        handle.swapAt(minNode, with: &removed)
+        handle.trickleDownMin(minNode)
+      }
+    }
+
+    _checkInvariants()
+    return removed
   }
 
   /// Removes and returns the element with the highest priority, if available.
   ///
   /// - Complexity: O(log `count`)
   @inlinable
   public mutating func popMax() -> Element? {
-    defer { _checkInvariants() }
-    guard count > 2 else {
-      // If count is 0, `popLast` will return `nil`
-      // If count is 1, the last (and only) item is the max
-      // If count is 2, the last item is the max (as it's the only item in the
-      // first max level)
-      return _storage.popLast()
-    }
-    // The max item is the larger of the two items in the first max level
-    let max = _storage.withUnsafeBufferPointer { buffer in
-      _Node(offset: buffer[2] >  buffer[1] ? 2 : 1, level: 1)
+    guard _storage.count > 2 else { return _storage.popLast() }
+
+    var removed = _storage.removeLast()
+
+    _update { handle in
+      if handle.count == 2 {
+        if handle.buffer[1] > removed {
+          handle.swapAt(.leftMax, with: &removed)
+        }
+      } else {
+        let maxNode = handle.buffer[2] > handle.buffer[1]
+          ? _Node.rightMax
+          : _Node.leftMax
+        handle.swapAt(maxNode, with: &removed)
+        handle.trickleDownMax(maxNode)
+      }
     }
-    return _remove(max)
+
+    _checkInvariants()
+    return removed
   }
 
   /// Removes and returns the element with the lowest priority.
@@ -149,30 +168,6 @@ public struct Heap<Element: Comparable> {
   public mutating func removeMax() -> Element {
     return popMax()!
   }
-
-  // MARK: -
-
-  @discardableResult
-  @inline(__always)
-  @inlinable
-  internal mutating func _remove(_ node: _Node) -> Element? {
-    guard _storage.count > node.offset else {
-      return nil
-    }
-
-    var removed = _storage.removeLast()
-
-    if node.offset < _storage.count {
-      _update { handle in
-        let p = handle.buffer.baseAddress.unsafelyUnwrapped
-        swap(&removed, &(p + node.offset).pointee)
-        //swap(&removed, &handle[node])
-        handle.trickleDown(node)
-      }
-    }
-
-    return removed
-  }
 }
 
 // MARK: -

Sources/PriorityQueueModule/_Node.swift

@@ -77,11 +77,24 @@ extension _Node {
 }
 
 extension _Node {
+  /// The root node in the heap.
   @inlinable @inline(__always)
   internal static var root: Self {
     Self.init(offset: 0, level: 0)
   }
 
+  /// The first max node in the heap. (I.e., the left child of the root.)
+  @inlinable @inline(__always)
+  internal static var leftMax: Self {
+    Self.init(offset: 1, level: 1)
+  }
+
+  /// The second max node in the heap. (I.e., the right child of the root.)
+  @inlinable @inline(__always)
+  internal static var rightMax: Self {
+    Self.init(offset: 2, level: 1)
+  }
+
   @inlinable
   internal var isMinLevel: Bool {
     Self.isMinLevel(level)