[HashTreeCollections] Add TreeDictionary.combining(_:by:)

Sources/HashTreeCollections/HashNode/_HashNode+Builder.swift

@@ -70,6 +70,16 @@ extension _HashNode.Builder {
     Self(level, .item(item, at: bucket))
   }
 
+  @inlinable @inline(__always)
+  internal static func anyNode(
+    _ level: _HashLevel, _ node: __owned _HashNode
+  ) -> Self {
+    if node.isCollisionNode {
+      return self.collisionNode(level, node)
+    }
+    return self.node(level, node)
+  }
+
   @inlinable @inline(__always)
   internal static func node(
     _ level: _HashLevel, _ node: __owned _HashNode
@@ -125,9 +135,33 @@ extension _HashNode.Builder {
     }
   }
 
+  @inlinable
+  internal init(
+    _ level: _HashLevel,
+    collisions1: __owned Self,
+    _ hash1: _Hash,
+    collisions2: __owned Self,
+    _ hash2: _Hash
+  ) {
+    assert(hash1 != hash2)
+    let b1 = hash1[level]
+    let b2 = hash2[level]
+    self = .empty(level)
+    if b1 == b2 {
+      let b = Self(
+        level.descend(),
+        collisions1: collisions1, hash1,
+        collisions2: collisions2, hash2)
+      self.addNewChildBranch(level, b, at: b1)
+    } else {
+      self.addNewChildBranch(level, collisions1, at: b1)
+      self.addNewChildBranch(level, collisions2, at: b2)
+    }
+  }
+
   @inlinable
   internal __consuming func finalize(_ level: _HashLevel) -> _HashNode {
-    assert(level.isAtRoot && self.level.isAtRoot)
+    //assert(level.isAtRoot && self.level.isAtRoot)
     switch kind {
     case .empty:
       return ._emptyNode()
@@ -193,10 +227,12 @@ extension _HashNode.Builder {
     _ level: _HashLevel, _ newItem: __owned Element, at newBucket: _Bucket
   ) {
     assert(level == self.level)
+    assert(!newBucket.isInvalid)
     switch kind {
     case .empty:
       kind = .item(newItem, at: newBucket)
     case .item(let oldItem, let oldBucket):
+      assert(!oldBucket.isInvalid)
       assert(oldBucket != newBucket)
       let node = _HashNode._regularNode(oldItem, oldBucket, newItem, newBucket)
       kind = .node(node)
@@ -215,6 +251,17 @@ extension _HashNode.Builder {
     }
   }
 
+  @inlinable
+  internal mutating func addNewItem(
+    _ level: _HashLevel,
+    _ key: Key,
+    _ value: __owned Value?,
+    at newBucket: _Bucket
+  ) {
+    guard let value = value else { return }
+    addNewItem(level, (key, value), at: newBucket)
+  }
+
   @inlinable
   internal mutating func addNewChildNode(
     _ level: _HashLevel, _ newChild: __owned _HashNode, at newBucket: _Bucket
@@ -358,3 +405,90 @@ extension _HashNode.Builder {
     return mapValues { _ in () }
   }
 }
+
+extension _HashNode.Builder {
+  @inlinable
+  internal static func conflictingItems(
+    _ level: _HashLevel,
+    _ item1: Element?,
+    _ item2: Element?,
+    at bucket: _Bucket
+  ) -> Self {
+    switch (item1, item2) {
+    case (nil, nil):
+      return .empty(level)
+    case let (item1?, nil):
+      return .item(level, item1, at: bucket)
+    case let (nil, item2?):
+      return .item(level, item2, at: bucket)
+    case let (item1?, item2?):
+      let h1 = _Hash(item1.key)
+      let h2 = _Hash(item2.key)
+      guard h1 != h2 else {
+        return .collisionNode(level, _HashNode._collisionNode(h1, item1, item2))
+      }
+      let n = _HashNode._build(
+        level: level.descend(),
+        item1: item1, h1,
+        item2: { $0.initialize(to: item2) }, h2)
+      return .node(level, n.top)
+    }
+  }
+
+  @inlinable
+  internal static func mergedUniqueBranch(
+    _ level: _HashLevel,
+    _ node: _HashNode,
+    by merge: (Element) throws -> Value?
+  ) rethrows -> Self {
+    try node.read { l in
+      var result = Self.empty(level)
+      if l.isCollisionNode {
+        let hash = l.collisionHash
+        for lslot: _HashSlot in .zero ..< l.itemsEndSlot {
+          let lp = l.itemPtr(at: lslot)
+          if let v = try merge(lp.pointee) {
+            result.addNewCollision(level, (lp.pointee.key, v), hash)
+          }
+        }
+        return result
+      }
+      for (bucket, lslot) in l.itemMap {
+        let lp = l.itemPtr(at: lslot)
+        let v = try merge(lp.pointee)
+        if let v = v {
+          result.addNewItem(level, (lp.pointee.key, v), at: bucket)
+        }
+      }
+      for (bucket, lslot) in l.childMap {
+        let b = try Self.mergedUniqueBranch(
+          level.descend(), l[child: lslot], by: merge)
+        result.addNewChildBranch(level, b, at: bucket)
+      }
+      return result
+    }
+  }
+
+  @inlinable
+  internal mutating func addNewItems(
+    _ level: _HashLevel,
+    at bucket: _Bucket,
+    item1: Element?,
+    item2: Element?
+  ) {
+    switch (item1, item2) {
+    case (nil, nil):
+      break
+    case let (item1?, nil):
+      self.addNewItem(level, item1, at: bucket)
+    case let (nil, item2?):
+      self.addNewItem(level, item2, at: bucket)
+    case let (item1?, item2?):
+      let n = _HashNode._build(
+        level: level,
+        item1: item1, _Hash(item1.key),
+        item2: { $0.initialize(to: item2) }, _Hash(item2.key))
+      self.addNewChildNode(level, n.top, at: bucket)
+    }
+  }
+}

Sources/HashTreeCollections/HashNode/_HashNode+Lookups.swift

@@ -26,8 +26,9 @@ extension _HashNode.UnsafeHandle {
     _ level: _HashLevel, _ key: Key, _ hash: _Hash
   ) -> (descend: Bool, slot: _HashSlot)? {
     guard !isCollisionNode else {
-      let r = _findInCollision(level, key, hash)
-      guard r.code == 0 else { return nil }
+      guard hash == collisionHash else { return nil }
+      let r = _findInCollision(key)
+      guard r.found else { return nil }
       return (false, r.slot)
     }
     let bucket = hash[level]
@@ -44,17 +45,12 @@ extension _HashNode.UnsafeHandle {
   }
 
   @inlinable @inline(never)
-  internal func _findInCollision(
-    _ level: _HashLevel, _ key: Key, _ hash: _Hash
-  ) -> (code: Int, slot: _HashSlot) {
+  internal func _findInCollision(_ key: Key) -> (found: Bool, slot: _HashSlot) {
     assert(isCollisionNode)
-    if !level.isAtBottom {
-      if hash != self.collisionHash { return (2, .zero) }
-    }
     // Note: this searches the items in reverse insertion order.
     guard let slot = reverseItems.firstIndex(where: { $0.key == key })
-    else { return (1, self.itemsEndSlot) }
-    return (0, _HashSlot(itemCount &- 1 &- slot))
+    else { return (false, self.itemsEndSlot) }
+    return (true, _HashSlot(itemCount &- 1 &- slot))
   }
 }
 
@@ -143,15 +139,15 @@ extension _HashNode.UnsafeHandle {
     _ level: _HashLevel, _ key: Key, _ hash: _Hash
   ) -> _FindResult {
     guard !isCollisionNode else {
-      let r = _findInCollision(level, key, hash)
-      if r.code == 0 {
-        return .found(.invalid, r.slot)
+      if hash != self.collisionHash {
+        assert(!level.isAtBottom)
+        return .expansion
       }
-      if r.code == 1 {
+      let r = _findInCollision(key)
+      guard r.found else {
         return .appendCollision
       }
-      assert(r.code == 2)
-      return .expansion
+      return .found(.invalid, r.slot)
     }
     let bucket = hash[level]
     if itemMap.contains(bucket) {

Sources/HashTreeCollections/HashNode/_HashNode+Structural subtracting.swift

@@ -124,7 +124,7 @@ extension _HashNode {
           var removing = false
 
           let ritems = r.reverseItems
-          for lslot: _HashSlot in stride(from: .zero, to: l.itemsEndSlot, by: 1) {
+          for lslot: _HashSlot in .zero ..< l.itemsEndSlot {
             let lp = l.itemPtr(at: lslot)
             let include = !ritems.contains { $0.key == lp.pointee.key }
             if include, removing {

Sources/HashTreeCollections/HashNode/_HashNode+Subtree Removals.swift

@@ -64,6 +64,34 @@ extension _HashNode {
   }
 }
 
+extension _HashNode {
+  @inlinable
+  internal func removing(
+    _ level: _HashLevel, _ bucket: _Bucket
+  ) -> (removed: Builder, replacement: Builder) {
+    read { handle in
+      assert(!handle.isCollisionNode)
+      if handle.itemMap.contains(bucket) {
+        let slot = handle.itemMap.slot(of: bucket)
+        let p = handle.itemPtr(at: slot)
+        let hash = _Hash(p.pointee.key)
+        let r = self.removing(level, p.pointee.key, hash)!
+        return (.item(level, r.removed, at: bucket), r.replacement)
+      } else if handle.childMap.contains(bucket) {
+        let slot = handle.childMap.slot(of: bucket)
+        if hasSingletonChild {
+          return (.anyNode(level.descend(), handle[child: slot]), .empty(level))
+        }
+        var remainder = self.copy()
+        let removed = remainder.removeChild(at: bucket, slot)
+        return (.anyNode(level.descend(), removed), .node(level, remainder))
+      } else {
+        return (.empty(level), .node(level, self))
+      }
+    }
+  }
+}
+
 extension _HashNode {
   @inlinable
   internal mutating func remove(