[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(

Sources/HashTreeCollections/TreeDictionary/TreeDictionary+Combine.swift

@@ -0,0 +1,63 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift Collections open source project
+//
+// Copyright (c) 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+//
+//===----------------------------------------------------------------------===//
+
+#if swift(>=5.7)
+@frozen
+public enum CombiningBehavior {
+  case include
+  case discard
+  case merge
+}
+
+public protocol TreeDictionaryCombiningStrategy<Key, Value> {
+  associatedtype Key: Hashable
+  associatedtype Value
+
+  var valuesOnlyInFirst: CombiningBehavior { get }
+  var valuesOnlyInSecond: CombiningBehavior { get }
+  var equalValuesInBoth: CombiningBehavior { get }
+  var unequalValuesInBoth: CombiningBehavior { get }
+
+  func areEquivalentValues(_ a: Value, _ b: Value) -> Bool
+
+  func merge(_ key: Key, _ value1: Value?, _ value2: Value?) throws -> Value?
+}
+
+extension TreeDictionaryCombiningStrategy {
+  public typealias Element = (key: Key, value: Value)
+}
+
+extension TreeDictionaryCombiningStrategy where Value: Equatable {
+  @inlinable @inline(__always)
+  public func areEquivalentValues(_ a: Value, _ b: Value) -> Bool {
+    a == b
+  }
+}
+
+extension TreeDictionary {
+  @inlinable
+  public func combining(
+    _ other: Self,
+    by strategy: some TreeDictionaryCombiningStrategy<Key, Value>
+  ) throws -> Self {
+    let root = try _root.combining(.top, other._root, by: strategy)
+    return Self(_new: root)
+  }
+
+  @inlinable
+  mutating func combine(
+    _ other: Self,
+    by strategy: some TreeDictionaryCombiningStrategy<Key, Value>
+  ) throws {
+    self = try combining(other, by: strategy)
+  }
+}
+#endif

Tests/HashTreeCollectionsTests/TreeDictionary Smoke Tests.swift

@@ -625,4 +625,60 @@ final class TreeDictionarySmokeTests: CollectionTestCase {
 
     expectTrue(expectedPositions.isEmpty)
   }
+
+  func test_combine() {
+    var d1 = TreeDictionary<Int, String>(uniqueKeysWithValues: (0 ..< 10000).map { ($0, "1") })
+    d1[1] = "1"
+    var d2 = d1
+//    for i in 10 ..< 20 {
+//      d1[i] = nil
+//    }
+//    for i in 20 ..< 30 {
+//      d2[i] = nil
+//    }
+    for i in 40 ..< 50 {
+      d2[i] = "2"
+    }
+
+    class TestStrategy: TreeDictionaryCombiningStrategy {
+      typealias Key = Int
+      typealias Value = String
+
+      var _equalCounter = 0
+      var _mergeCounter = 0
+
+      var valuesOnlyInFirst: CombiningBehavior { .merge }
+      var valuesOnlyInSecond: CombiningBehavior { .merge }
+      var equalValuesInBoth: CombiningBehavior { .discard }
+      var unequalValuesInBoth: CombiningBehavior { .merge }
+
+      func areEquivalentValues(_ a: Value, _ b: Value) -> Bool {
+        _equalCounter += 1
+        return a == b
+      }
+
+      func merge(
+        _ key: Key, _ value1: Value?, _ value2: Value?
+      ) throws -> Value? {
+        _mergeCounter += 1
+
+        let s1 = value1 ?? "nil"
+        let s2 = value2 ?? "nil"
+        print("key: \(key), value1: \(s1), value2: \(s2)")
+
+        switch (value1, value2) {
+        case (nil, nil): return "00"
+        case (_?, nil): return "10"
+        case (nil, _?): return "01"
+        case (_?, _?): return "11"
+        }
+      }
+    }
+
+    let strategy = TestStrategy()
+    let d = try! d1.combining(d2, by: strategy)
+    print(d.map { ($0.key, $0.value) }.sorted(by: { $0.0 < $1.0 }))
+    print("Merge count: \(strategy._mergeCounter)")
+    print("isEqual count: \(strategy._equalCounter)")
+  }
 }

Utils/swift-collections.xcworkspace/xcshareddata/xcschemes/HashTreeCollections.xcscheme

@@ -20,6 +20,20 @@
                ReferencedContainer = "container:..">
             </BuildableReference>
          </BuildActionEntry>
+         <BuildActionEntry
+            buildForTesting = "YES"
+            buildForRunning = "YES"
+            buildForProfiling = "NO"
+            buildForArchiving = "NO"
+            buildForAnalyzing = "NO">
+            <BuildableReference
+               BuildableIdentifier = "primary"
+               BlueprintIdentifier = "HashTreeCollectionsTests"
+               BuildableName = "HashTreeCollectionsTests"
+               BlueprintName = "HashTreeCollectionsTests"
+               ReferencedContainer = "container:..">
+            </BuildableReference>
+         </BuildActionEntry>
       </BuildActionEntries>
    </BuildAction>
    <TestAction