Persistent collections updates (part 7)

Benchmarks/Sources/CppBenchmarks/src/MapBenchmarks.cpp

@@ -12,7 +12,7 @@
 #include "MapBenchmarks.h"
 #include <map>
 #include <cstdlib>
-#include "utils.h"
+#include "Utils.h"
 
 typedef std::map<intptr_t, intptr_t> custom_map;
 

Package.swift

@@ -96,9 +96,10 @@ let package = Package(
 
     // BitSet, BitArray
     .target(
-        name: "BitCollections",
-        path: "Sources/BitCollections",
-        swiftSettings: settings),
+      name: "BitCollections",
+      dependencies: ["_CollectionsUtilities"],
+      path: "Sources/BitCollections",
+      swiftSettings: settings),
     .testTarget(
       name: "BitCollectionsTests",
       dependencies: ["BitCollections", "_CollectionsTestSupport"],
@@ -118,6 +119,7 @@ let package = Package(
     // Heap<Value>
     .target(
         name: "PriorityQueueModule",
+        dependencies: ["_CollectionsUtilities"],
         exclude: ["CMakeLists.txt"],
         swiftSettings: settings),
     .testTarget(
@@ -149,6 +151,7 @@ let package = Package(
     // SortedSet<Element>, SortedDictionary<Key, Value>
     .target(
       name: "SortedCollections",
+      dependencies: ["_CollectionsUtilities"],
       swiftSettings: settings),
     .testTarget(
       name: "SortedCollectionsTests",

Sources/BitCollections/BitSet/BitSet.Counted.swift

@@ -29,9 +29,8 @@ extension BitSet.Counted {
   @inline(never)
   @_effects(releasenone)
   public func _checkInvariants() {
-    precondition(_count == _bits.count, "Invalid count")
-    precondition(_storage.isEmpty || !_storage.last!.isEmpty,
-                 "Extraneous tail slot")
+    _bits._checkInvariants()
+    precondition(_count == _bits.count)
   }
 #else
   @inline(__always) @inlinable

Sources/PersistentCollections/Node/_AncestorOffsets.swift → Sources/PersistentCollections/Node/_AncestorSlots.swift

@@ -52,7 +52,7 @@ extension _AncestorSlots {
       return _Slot((path &>> level.shift) & _Bucket.bitMask)
     }
     set {
-      assert(newValue._value < UInt.bitWidth)
+      assert(newValue._value < _Bitmap.capacity)
       assert(self[level] == .zero)
       path |= (UInt(truncatingIfNeeded: newValue._value) &<< level.shift)
     }

Sources/PersistentCollections/Node/_Bitmap.swift

@@ -116,9 +116,14 @@ extension _Bitmap {
 }
 
 extension _Bitmap {
+  @inlinable @inline(__always)
+  internal func isSubset(of other: Self) -> Bool {
+    _value & ~other._value == 0
+  }
+
   @inlinable @inline(__always)
   internal func isDisjoint(with other: Self) -> Bool {
-    _value & other._value != 0
+    _value & other._value == 0
   }
 
   @inlinable @inline(__always)

Sources/PersistentCollections/Node/_HashTreeIterator.swift

@@ -0,0 +1,129 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+@usableFromInline
+@frozen
+internal struct _HashTreeIterator {
+  @usableFromInline
+  internal struct _Opaque {
+    internal var ancestorSlots: _AncestorSlots
+    internal var ancestorNodes: _Stack<_UnmanagedNode>
+    internal var level: _Level
+    internal var isAtEnd: Bool
+
+    @usableFromInline
+    @_effects(releasenone)
+    internal init(_ root: _UnmanagedNode) {
+      self.ancestorSlots = .empty
+      self.ancestorNodes = _Stack(filledWith: root)
+      self.level = .top
+      self.isAtEnd = false
+    }
+  }
+
+  @usableFromInline
+  internal let root: _RawStorage
+
+  @usableFromInline
+  internal var node: _UnmanagedNode
+
+  @usableFromInline
+  internal var slot: _Slot
+
+  @usableFromInline
+  internal var endSlot: _Slot
+
+  @usableFromInline
+  internal var _o: _Opaque
+
+  @usableFromInline
+  @_effects(releasenone)
+  internal init(root: __shared _RawNode) {
+    self.root = root.storage
+    self.node = root.unmanaged
+    self.slot = .zero
+    self.endSlot = node.itemsEndSlot
+    self._o = _Opaque(self.node)
+
+    if node.hasItems { return }
+    if node.hasChildren {
+      _descendToLeftmostItem(ofChildAtSlot: .zero)
+    } else {
+      self._o.isAtEnd = true
+    }
+  }
+}
+
+extension _HashTreeIterator: IteratorProtocol {
+  @inlinable
+  internal mutating func next() -> (node: _UnmanagedNode, slot: _Slot)? {
+    guard slot < endSlot else {
+      return _next()
+    }
+    defer { slot = slot.next() }
+    return (node, slot)
+  }
+
+  @usableFromInline
+  @_effects(releasenone)
+  internal mutating func _next() -> (node: _UnmanagedNode, slot: _Slot)? {
+    if _o.isAtEnd { return nil }
+    if node.hasChildren {
+      _descendToLeftmostItem(ofChildAtSlot: .zero)
+      slot = slot.next()
+      return (node, .zero)
+    }
+    while !_o.level.isAtRoot {
+      let nextChild = _ascend().next()
+      if nextChild < node.childrenEndSlot {
+        _descendToLeftmostItem(ofChildAtSlot: nextChild)
+        slot = slot.next()
+        return (node, .zero)
+      }
+    }
+    // At end
+    endSlot = node.itemsEndSlot
+    slot = endSlot
+    _o.isAtEnd = true
+    return nil
+  }
+}
+
+extension _HashTreeIterator {
+  internal mutating func _descend(toChildSlot childSlot: _Slot) {
+    assert(childSlot < node.childrenEndSlot)
+    _o.ancestorSlots[_o.level] = childSlot
+    _o.ancestorNodes.push(node)
+    _o.level = _o.level.descend()
+    node = node.unmanagedChild(at: childSlot)
+    slot = .zero
+    endSlot = node.itemsEndSlot
+  }
+
+  internal mutating func _ascend() -> _Slot {
+    assert(!_o.level.isAtRoot)
+    node = _o.ancestorNodes.pop()
+    _o.level = _o.level.ascend()
+    let childSlot = _o.ancestorSlots[_o.level]
+    _o.ancestorSlots.clear(_o.level)
+    return childSlot
+  }
+
+  internal mutating func _descendToLeftmostItem(
+    ofChildAtSlot childSlot: _Slot
+  ) {
+    _descend(toChildSlot: childSlot)
+    while endSlot == .zero {
+      assert(node.hasChildren)
+      _descend(toChildSlot: .zero)
+    }
+  }
+}

Sources/PersistentCollections/Node/_Level.swift

@@ -32,6 +32,12 @@ internal struct _Level {
     assert(shift <= UInt8.max)
     self._shift = UInt8(truncatingIfNeeded: shift)
   }
+
+  @inlinable @inline(__always)
+  init(depth: Int) {
+    assert(depth > 0 && depth < _Level.limit)
+    self.init(shift: UInt(bitPattern: depth * _Bitmap.bitWidth))
+  }
 }
 
 extension _Level {

Sources/PersistentCollections/Node/_Node+Invariants.swift

@@ -55,7 +55,7 @@ extension _Node {
       if $0.isCollisionNode {
         precondition(count == $0.itemCount)
         precondition(count > 0)
-        let key = $0[item: .zero].key
+        let key = $0.collisionHash
         let hash = _Hash(key)
         precondition(
           hash.isEqual(to: path, upTo: level),

Sources/PersistentCollections/Node/_Node+Lookups.swift

@@ -49,7 +49,7 @@ extension _Node.UnsafeHandle {
   ) -> (code: Int, slot: _Slot, expansionHash: _Hash) {
     assert(isCollisionNode)
     if !level.isAtBottom {
-      let h = _Hash(self[item: .zero].key)
+      let h = self.collisionHash
       if h != hash { return (2, .zero, h) }
     }
     // Note: this searches the items in reverse insertion order.
@@ -182,21 +182,30 @@ extension _Node {
       }
     }
   }
+}
 
+extension _Node {
   @inlinable
   internal func containsKey(
     _ level: _Level, _ key: Key, _ hash: _Hash
   ) -> Bool {
-    read {
-      let r = $0.find(level, key, hash, forInsert: false)
-      switch r {
-      case .found:
-        return true
-      case .notFound, .newCollision, .expansion:
-        return false
-      case .descend(_, let slot):
-        return $0[child: slot].containsKey(level.descend(), key, hash)
-      }
+    read { $0.containsKey(level, key, hash) }
+  }
+}
+
+extension _Node.UnsafeHandle {
+  @inlinable
+  internal func containsKey(
+    _ level: _Level, _ key: Key, _ hash: _Hash
+  ) -> Bool {
+    let r = find(level, key, hash, forInsert: false)
+    switch r {
+    case .found:
+      return true
+    case .notFound, .newCollision, .expansion:
+      return false
+    case .descend(_, let slot):
+      return self[child: slot].containsKey(level.descend(), key, hash)
     }
   }
 }