Documentation refactors, polishes, and add SubSequences

Key Changes:
 - Add custom SubSequences
 - Custom SortedDictionary.(Keys|Values).Sequence iterator
 - Add SortedDictionary.modifyValue(forKey:default:_:)

Minor Changes:
 - Update documentation
 - Various API optimizations
 - Fix bug where offsetBy:limitedBy: calls did not validate indices
 - Remove unique initializers with documented duplicate handling behavior.
 - Implement _Node.UnsafeHandle.append*() methods
 - SortedSet no longer allocates a 'values' buffer (_Node.Header.values is now optional)

Documentation:
 - Complexity guarantees
   - BidirectionalCollection conformances
   - Sequence/IteratorProtocol methods

Optimizations:
 - Add optimized dictionary value mapping (through Node.init(mappingFrom:_:))
 - Custom _BTree.forEach
 - Greater explicit @inline(__always) annotations throughout.
 - Simplify _BTree: Sequence implementation
 - Move methods to using _BTree.Builder
 - _BTree.FixedSizeArray now uses `withUnsafePointer` to prevent branching.

Author: vihanb

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

@@ -20,11 +20,11 @@ extension _BTree {
     _ isIncluded: (Element) throws -> Bool
   ) rethrows -> _BTree {
     // TODO: optimize implementation to O(n)
-    var newTree: _BTree = _BTree()
+    var builder = Builder()
     for element in self where try isIncluded(element) {
-      newTree.updateAnyValue(element.value, forKey: element.key)
+      builder.append(element)
     }
-    return newTree
+    return builder.finish()
   }
 
 

Sources/SortedCollections/BTree/_BTree+Sequence.swift

@@ -10,113 +10,166 @@
 //===----------------------------------------------------------------------===//
 
 extension _BTree: Sequence {
+  @inlinable
+  internal func forEach(_ body: (Element) throws -> Void) rethrows {
+    func loop(node: Unmanaged<Node.Storage>) throws {
+      try node._withUnsafeGuaranteedRef { storage in
+        try storage.read { handle in
+          for i in 0..<handle.elementCount {
+            if !handle.isLeaf {
+              try loop(node: .passUnretained(handle[childAt: i].storage))
+            }
+            
+            try body(handle[elementAt: i])
+          }
+          
+          if !handle.isLeaf {
+            let lastChild: Unmanaged =
+              .passUnretained(handle[childAt: handle.childCount - 1].storage)
+            try loop(node: lastChild)
+          }
+        }
+      }
+    }
+    
+    try loop(node: .passUnretained(self.root.storage))
+  }
+  
   @usableFromInline
   internal struct Iterator: IteratorProtocol {
     @usableFromInline
     internal let tree: _BTree
 
     @usableFromInline
-    internal var offsets: Index.Offsets
-    
-    @usableFromInline
-    internal var parents: FixedSizeArray<Unmanaged<Node.Storage>>
+    internal var slots: FixedSizeArray<Slot>
 
     @usableFromInline
-    internal var currentNode: Unmanaged<Node.Storage>
+    internal var path: FixedSizeArray<Unmanaged<Node.Storage>>
 
-    @usableFromInline
-    internal var slot: Int
+    /// Creates an iterator to the element within a tree corresponding to a specific index
+    @inlinable
+    @inline(__always)
+    internal init(forTree tree: _BTree, startingAt index: Index) {
+      self.tree = tree
+      self.slots = index.childSlots
+      self.slots.append(UInt16(index.slot))
+      
+      self.path = .init(repeating: .passUnretained(tree.root.storage))
+      
+      if self.tree.isEmpty {
+        return
+      }
+      
+      var node: Unmanaged = .passUnretained(tree.root.storage)
+      for depth in 0..<index.childSlots.depth {
+        let childSlot = index.childSlots[depth]
+        self.path.append(node)
+        
+        node._withUnsafeGuaranteedRef {
+          $0.read { handle in
+            node = .passUnretained(handle[childAt: Int(childSlot)].storage)
+          }
+        }
+      }
+    }
 
+    /// Creates an iterator to the first element within a tree.
     @inlinable
     @inline(__always)
-    internal init(tree: _BTree) {
+    internal init(forTree tree: _BTree) {
       self.tree = tree
 
-      self.offsets = .init(repeating: 0)
-      self.parents = .init(repeating: .passUnretained(tree.root.storage))
+      self.slots = .init(repeating: 0)
+      self.path = .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 }) {
-        self.parents.append(Unmanaged.passUnretained(node.storage))
-        self.offsets.append(0)
+      var nextNode: Unmanaged? = .passUnretained(tree.root.storage)
+      while let node = nextNode {
+        self.path.append(node)
+        self.slots.append(0)
 
-        node = node.read({ $0[childAt: 0] })
+        node._withUnsafeGuaranteedRef {
+          $0.read { handle in
+            if handle.isLeaf {
+              nextNode = nil
+            } else {
+              nextNode = .passUnretained(handle[childAt: 0].storage)
+            }
+          }
+        }
       }
-      
-
-      self.currentNode = .passUnretained(node.storage)
-      self.slot = 0
     }
 
     @inlinable
     @inline(__always)
     internal mutating func next() -> Element? {
       // Check slot sentinel value for end of tree.
-      if _slowPath(self.slot == -1) {
+      if _slowPath(path.isEmpty) {
         return nil
       }
 
-      return self.currentNode._withUnsafeGuaranteedRef { storage in
+      return path.last._withUnsafeGuaranteedRef { storage in
         storage.read({ handle in
           defer {
             // If we're not a leaf, descend to the next child
             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]
+              // Go to the right child
+              slots.last += 1
+              var nextNode: Unmanaged? = .passUnretained(handle[childAt: Int(slots.last)].storage)
 
-              while !node.read({ $0.isLeaf }) {
-                self.parents.append(.passUnretained(node.storage))
-                self.offsets.append(0)
-                node = node.read({ $0[childAt: 0] })
+              while let node = nextNode {
+                self.path.append(node)
+                self.slots.append(0)
+                
+                node._withUnsafeGuaranteedRef {
+                  $0.read { handle in
+                    if handle.isLeaf {
+                      nextNode = nil
+                    } else {
+                      nextNode = .passUnretained(handle[childAt: 0].storage)
+                    }
+                  }
+                }
               }
-              
-              self.currentNode = .passUnretained(node.storage)
-              self.slot = 0
             } else {
-              if self.slot < handle.elementCount - 1 {
-                self.slot += 1
+              if slots.last < handle.elementCount - 1 {
+                slots.last += 1
               } else {
-                while true {
-                  if self.parents.depth == 0 {
-                    self.slot = -1
-                    break
-                  }
+                _ = path.pop()
+                _ = slots.pop()
+                
+                while !path.isEmpty {
+                  let parent = path.last
+                  let slot = slots.last
 
-                  // If we are at a leaf, then ascend to the
-                  let parent = self.parents.pop()
-                  let offset = self.offsets.pop()
-                  
-                  let parentElements = parent._withUnsafeGuaranteedRef { $0.read({ $0.elementCount }) }
+                  let parentElementCount = parent._withUnsafeGuaranteedRef {
+                    $0.read({ $0.elementCount })
+                  }
 
-                  if offset < parentElements {
-                    self.currentNode = parent
-                    self.slot = Int(offset)
+                  if slot < parentElementCount {
                     break
+                  } else {
+                    _ = path.pop()
+                    _ = slots.pop()
                   }
                 }
               }
             }
           }
 
-          return storage.read({ $0[elementAtSlot: self.slot] })
+          return storage.read({ $0[elementAt: Int(slots.last)] })
         })
       }
     }
   }
 
   @inlinable
   internal func makeIterator() -> Iterator {
-    return Iterator(tree: self)
+    return Iterator(forTree: self)
   }
 }