Continue applying refactors and optimizations.

Key Changes:
 - _BTree.init() dynamically computes capacity from stride
 - _BTree.UnsafePath uses Unmanaged +0 references.

Minor Changes:
 - Fix _BTree.Iterator bug on single-leaf trees.

Author: vihanb

Sources/SortedCollections/BTree/_BTree+BidirectionalCollection.swift

@@ -14,7 +14,7 @@ extension _BTree: BidirectionalCollection {
   /// - Complexity: O(1)
   @inlinable
   @inline(__always)
-  internal var count: Int { self.root.storage.header.totalElements }
+  internal var count: Int { self.root.storage.header.subtreeCount }
 
   /// A Boolean value that indicates whether the BTree is empty.
   @inlinable
@@ -61,7 +61,8 @@ extension _BTree: BidirectionalCollection {
   /// - Parameters:
   ///   - start: A valid index of the collection.
   ///   - end: Another valid index of the collection. If end is equal to start, the result is zero.
-  /// - Returns: The distance between start and end. The result can be negative only if the collection conforms to the BidirectionalCollection protocol.
+  /// - Returns: The distance between start and end. The result can be negative only if the collection
+  ///     conforms to the BidirectionalCollection protocol.
   /// - Complexity: O(1)
   @inlinable
   internal func distance(from start: Index, to end: Index) -> Int {
@@ -77,9 +78,9 @@ extension _BTree: BidirectionalCollection {
       preconditionFailure("Attempt to advance out of collection bounds.")
     }
 
-    let shouldSeekWithinLeaf = Node(path.node).read({
+    let shouldSeekWithinLeaf = path.readNode {
       $0.isLeaf && _fastPath(path.slot + 1 < $0.elementCount)
-    })
+    }
 
     if shouldSeekWithinLeaf {
       // Continue searching within the same leaf
@@ -143,10 +144,10 @@ extension _BTree: BidirectionalCollection {
 
     // TODO: optimization for searching within children
 
-    if var path = i.path, path.node.header.children == nil {
+    if var path = i.path, path.readNode({ $0.isLeaf }) {
       // Check if the target element will be in the same node
       let targetSlot = path.slot + distance
-      if 0 <= targetSlot && targetSlot < path.node.header.count {
+      if 0 <= targetSlot && targetSlot < path.readNode({ $0.elementCount }) {
         path.slot = targetSlot
         path.offset = newIndex
         i.path = path

Sources/SortedCollections/BTree/_BTree+Partial RangeReplaceableCollection.swift

@@ -22,7 +22,7 @@ extension _BTree {
     // TODO: optimize implementation to O(n)
     var newTree: _BTree = _BTree()
     for element in self where try isIncluded(element) {
-      newTree.setAnyValue(element.value, forKey: element.key)
+      newTree.updateAnyValue(element.value, forKey: element.key)
     }
     return newTree
   }
@@ -102,21 +102,21 @@ extension _BTree {
     }
   }
 
-  // MARK: Offset Removal
-  /// Removes the element of a tree at a given offset.
+  // MARK: Index Removal
+  /// Removes the element of a tree at a given index.
   ///
-  /// - Parameter offset: the offset which must be in-bounds.
+  /// - Parameter index: a valid index of the tree, not `endIndex`
   /// - Returns: The moved element of the tree
   @inlinable
   @inline(__always)
   @discardableResult
-  internal mutating func remove(at offset: Int) -> Element {
+  internal mutating func remove(at index: Index) -> Element {
     invalidateIndices()
-    let removedElement = self.root.update { $0.remove(at: offset) }
-    self._balanceRoot()
-    return removedElement
+    guard let path = index.path else { preconditionFailure("Index out of bounds.") }
+    return self.remove(atOffset: path.offset)
   }
 
+  // MARK: Bulk Removal
   @inlinable
   @inline(__always)
   internal mutating func removeAll() {
@@ -125,20 +125,6 @@ extension _BTree {
     self.root = _Node(withCapacity: _BTree.defaultLeafCapacity, isLeaf: true)
   }
 
-  // MARK: Index Removal
-  /// Removes the element of a tree at a given index.
-  ///
-  /// - Parameter index: a valid index of the tree, not `endIndex`
-  /// - Returns: The moved element of the tree
-  @inlinable
-  @inline(__always)
-  @discardableResult
-  internal mutating func remove(at index: Index) -> Element {
-    invalidateIndices()
-    guard let path = index.path else { preconditionFailure("Index out of bounds.") }
-    return self.remove(at: path.offset)
-  }
-  
   /// Removes the elements in the specified subrange from the collection.
   @inlinable
   internal mutating func removeSubrange(_ bounds: Range<Index>) {
@@ -149,7 +135,7 @@ extension _BTree {
     let startOffset = startPath.offset
 
     for _ in 0..<rangeSize {
-      self.remove(at: startOffset)
+      self.remove(atOffset: startOffset)
     }
   }
 

Sources/SortedCollections/BTree/_BTree+Sequence.swift

@@ -35,6 +35,13 @@ extension _BTree: Sequence {
       self.offsets = .init(repeating: 0)
       self.parents = .init(repeating: .passUnretained(tree.root.storage))
 
+      // Simple case for an empty tree
+      if self.tree.isEmpty {
+        self.currentNode = .passUnretained(tree.root.storage)
+        self.slot = -1
+        return
+      }
+      
       // TODO: maybe convert to unowned(unsafe)
       var node = tree.root
       while !node.read({ $0.isLeaf }) {
@@ -52,6 +59,7 @@ extension _BTree: Sequence {
     @inlinable
     @inline(__always)
     internal mutating func next() -> Element? {
+      // Check slot sentinel value for end of tree.
       if _slowPath(self.slot == -1) {
         return nil
       }
@@ -63,11 +71,13 @@ extension _BTree: Sequence {
             if !handle.isLeaf {
               self.parents.append(self.currentNode)
               self.offsets.append(UInt16(self.slot + 1))
+              
+              // TODO: make these descents Unmanaged
               var node = handle[childAt: self.slot + 1]
 
               while !node.read({ $0.isLeaf }) {
-                self.offsets.append(0)
                 self.parents.append(.passUnretained(node.storage))
+                self.offsets.append(0)
                 node = node.read({ $0[childAt: 0] })
               }
 
@@ -78,6 +88,11 @@ extension _BTree: Sequence {
                 self.slot += 1
               } else {
                 while true {
+                  if self.parents.depth == 0 {
+                    self.slot = -1
+                    break
+                  }
+                  
                   // If we are at a leaf, then ascend to the
                   let parent = self.parents.pop()
                   let offset = self.offsets.pop()
@@ -88,9 +103,6 @@ extension _BTree: Sequence {
                     self.currentNode = parent
                     self.slot = Int(offset)
                     break
-                  } else if self.parents.depth == 0 {
-                    self.slot = -1
-                    break
                   }
                 }
               }

Sources/SortedCollections/BTree/_BTree.Index.swift

@@ -41,7 +41,7 @@ extension _BTree {
     @inlinable
     @inline(__always)
     internal static func ==(lhs: Index, rhs: Index) -> Bool {
-      return lhs.path?.node === rhs.path?.node
+      return lhs.path?.offset == rhs.path?.offset
     }
 
     @inlinable
@@ -73,7 +73,9 @@ extension _BTree {
     @inline(__always)
     internal func ensureValid(with index: Index) {
       precondition(
-        self.root != nil && self.root === index.root && self.version == index.version,
+        self.root != nil &&
+          self.root === index.root &&
+          self.version == index.version,
         "Attempt to use an invalid indices.")
     }
   }

Sources/SortedCollections/BTree/_BTree.PackedList.swift

@@ -9,15 +9,17 @@
 //
 //===----------------------------------------------------------------------===//
 
-@usableFromInline
-internal let _PACKED_OFFSET_LIST_MAX_SIZE = 16
-
 extension _BTree {
   /// A stack-allocated list of some values.
+  /// 
   /// - Warning: This may hold strong references to objects after they
   ///     do not put non-trivial types in this.
   @usableFromInline
   internal struct FixedSizeArray<Value> {
+    @inlinable
+    @inline(__always)
+    internal static var maxSize: UInt8 { 16 }
+    
     @usableFromInline
     internal var depth: UInt8
 
@@ -45,7 +47,8 @@ extension _BTree {
     @inlinable
     @inline(__always)
     internal mutating func append(_ offset: Value) {
-      assert(depth < _PACKED_OFFSET_LIST_MAX_SIZE, "Out of bounds access in offset list.")
+      assert(depth < FixedSizeArray.maxSize,
+             "Out of bounds access in offset list.")
       self.depth &+= 1
       self[self.depth] = offset
     }
@@ -63,7 +66,8 @@ extension _BTree {
     @inline(__always)
     internal subscript(_ offset: UInt8) -> Value {
       get {
-        assert(offset <= depth && depth <= _PACKED_OFFSET_LIST_MAX_SIZE, "Out of bounds access in offset list.")
+        assert(offset <= depth && depth <= FixedSizeArray.maxSize,
+               "Out of bounds access in offset list.")
 
         switch offset {
         case 0: return self.values.0
@@ -87,7 +91,9 @@ extension _BTree {
       }
 
       _modify {
-        assert(offset <= depth && depth <= _PACKED_OFFSET_LIST_MAX_SIZE, "Out of bounds access in offset list.")
+        assert(offset <= depth && depth <= FixedSizeArray.maxSize,
+               "Out of bounds access in offset list.")
+        
         switch offset {
         case 0: yield &self.values.0
         case 1: yield &self.values.1

Sources/SortedCollections/BTree/_BTree.UnsafePath.swift

@@ -28,9 +28,8 @@ extension _BTree {
     @usableFromInline
     internal var childSlots: Offsets
 
-    // TODO: switch to Unmanaged<Node.Storage>
     @usableFromInline
-    internal unowned(unsafe) var node: Node.Storage
+    internal var node: Unmanaged<Node.Storage>
 
     @usableFromInline
     internal var slot: Int
@@ -85,7 +84,7 @@ extension _BTree {
       childSlots: Offsets,
       offset: Int
     ) {
-      self.node = node
+      self.node = .passUnretained(node)
       self.slot = slot
       self.childSlots = childSlots
       self.offset = offset
@@ -97,7 +96,16 @@ extension _BTree {
     @inlinable
     @inline(__always)
     internal var element: Element {
-      return Node(self.node).read { $0[elementAt: self.slot] }
+      return self.readNode { $0[elementAt: self.slot] }
+    }
+    
+    /// Operators on a handle of the node
+    @inlinable
+    @inline(__always)
+    internal func readNode<R>(
+      _ body: (Node.UnsafeHandle) throws -> R
+    ) rethrows -> R {
+      return try self.node._withUnsafeGuaranteedRef { try $0.read(body) }
     }
   }
 }
@@ -110,7 +118,7 @@ extension _BTree.UnsafePath: Equatable {
   @inlinable
   public static func ==(lhs: _BTree.UnsafePath, rhs: _BTree.UnsafePath) -> Bool {
     // We assume the parents are the same
-    return lhs.node === rhs.node && lhs.slot == rhs.slot
+    return lhs.node.toOpaque() == rhs.node.toOpaque() && lhs.slot == rhs.slot
   }
 }