Initial implementation of SetAlgebra and various refactors.

Key Changes:
 - SortedSet
   - Implement SetAlgebra
   - Implement Hashable
   - Implement Equatable
   - Implement Custom(Debug)StringConvertible
   - Implement CustomReflectable
   - Implement Codable
   - Implement subscript(range: Range<Element>)
   - Implement subscript(offset:)
   - Implement remove(atOffset:)

Minor Changes:
 - Fix inefficiency on SortedDictionary.customMirror
 - Optimize to O(1) _BTree.UnsafePath comparator
 - Optimize away ARC traffic on findAnyPath(forKey:)

Various additional refactors and renames.

Author: vihanb

Sources/SortedCollections/BTree/_BTree+BidirectionalCollection.swift

@@ -21,8 +21,9 @@ extension _BTree: BidirectionalCollection {
   @inline(__always)
   internal var isEmpty: Bool { self.count == 0 }
 
+  // TODO: look into O(1) implementation
   /// Locates the first element and returns a proper path to it, or nil if the BTree is empty.
-  /// - Complexity: O(1)
+  /// - Complexity: O(`log n`)
   @inlinable
   internal var startIndex: Index {
     if count == 0 { return Index(nil, forTree: self) }
@@ -156,7 +157,7 @@ extension _BTree: BidirectionalCollection {
     }
 
     // Otherwise, re-seek
-    i = self.indexToElement(at: newIndex)
+    i = Index(self.path(atOffset: newIndex), forTree: self)
   }
 
   /// Returns an index that is the specified distance from the given index.
@@ -225,98 +226,126 @@ extension _BTree: BidirectionalCollection {
   @inline(__always)
   internal subscript(index: Index) -> Element {
     assert(index.path != nil, "Attempt to subscript out of range index.")
-    return index.path!.element
+    return index.path.unsafelyUnwrapped.element
   }
 }
 
 // MARK: Custom Indexing Operations
 extension _BTree {
-  /// Returns the path corresponding to the first found instance of the key. This may
-  /// not be the first instance of the key. This is marginally more efficient for trees
-  /// that do not have duplicates.
-  ///
-  /// - Parameter key: The key to search for within the tree.
-  /// - Returns: If found, returns a path to the element. Otherwise, `nil`.
+  /// Returns a path to the key at absolute offset `i`.
+  /// - Parameter offset: 0-indexed offset within BTree bounds, else may panic.
+  /// - Returns: the index of the appropriate element.
+  /// - Complexity: O(`log n`)
+  
+
+  /// Obtains the start index for a key (or where it would exist).
   @inlinable
-  internal func anyIndex(forKey key: Key) -> Index? {
+  internal func startIndex(forKey key: Key) -> Index {
     var childSlots = UnsafePath.Offsets(repeating: 0)
-    var node: Node? = self.root
+    var targetSlot: Int = 0
+    var offset = 0
 
-    while let currentNode = node {
-      let path: UnsafePath? = currentNode.read { handle in
-        let keySlot = handle.startSlot(forKey: key)
-        if keySlot < handle.elementCount && handle[keyAt: keySlot] == key {
-          return UnsafePath(node: currentNode.storage, slot: keySlot, childSlots: childSlots, offset: 0)
-        } else {
+    func search(in node: Node) -> Unmanaged<Node.Storage>? {
+      node.read({ handle in
+        let slot = handle.startSlot(forKey: key)
+        if slot < handle.elementCount {
           if handle.isLeaf {
-            node = nil
+            offset += slot
+            targetSlot = slot
+            return .passUnretained(node.storage)
           } else {
-            childSlots.append(UInt16(keySlot))
-            node = handle[childAt: keySlot]
+            // Calculate offset by summing previous subtrees
+            for i in 0...slot {
+              offset += handle[childAt: i].read({ $0.subtreeCount })
+            }
+            
+            let currentOffset = offset
+            let currentDepth = childSlots.depth
+            childSlots.append(UInt16(slot))
+            
+            if let foundEarlier = search(in: handle[childAt: slot]) {
+              return foundEarlier
+            } else {
+              childSlots.depth = currentDepth
+              targetSlot = slot
+              offset = currentOffset
+              
+              return .passUnretained(node.storage)
+            }
           }
-          
+        } else {
+          // Start index exceeds node and is therefore not in this.
           return nil
         }
-      }
-      
-      if let path = path {
-        return Index(path, forTree: self)
-      }
+      })
     }
 
-    return nil
+    if let targetChild = search(in: self.root) {
+      return Index(UnsafePath(
+        node: targetChild,
+        slot: targetSlot,
+        childSlots: childSlots,
+        offset: offset
+      ), forTree: self)
+    } else {
+      return Index(nil, forTree: self)
+    }
   }
 
-  /// Returns a path to the key at absolute offset `i`.
-  /// - Parameter offset: 0-indexed offset within BTree bounds, else may panic.
-  /// - Returns: the index of the appropriate element.
-  /// - Complexity: O(`log n`)
+  /// Obtains the last index at which a value less than or equal to the key appears.
   @inlinable
-  internal func indexToElement(at offset: Int) -> Index {
-    assert(offset <= self.count, "Index out of bounds.")
-    
-    if offset == self.count {
-      return Index(nil, forTree: self)
-    }
-    
+  internal func lastIndex(forKey key: Key) -> Index {
     var childSlots = UnsafePath.Offsets(repeating: 0)
+    var targetSlot: Int = 0
+    var offset = 0
 
-    var node: _Node = self.root
-    var startIndex = 0
-    
-    while !node.read({ $0.isLeaf }) {
-      let internalPath: UnsafePath? = node.read { handle in
-        for childSlot in 0..<handle.childCount {
-          let child = handle[childAt: childSlot]
-          let endIndex = startIndex + child.read({ $0.subtreeCount })
+    func search(in node: Node) -> Unmanaged<Node.Storage>? {
+      node.read({ handle in
+        let slot = handle.endSlot(forKey: key) - 1
+        if slot > 0 {
+          // Sanity Check
+          assert(slot < handle.elementCount, "Slot out of bounds.")
 
-          if offset < endIndex {
-            childSlots.append(UInt16(childSlot))
-            node = child
-            return nil
-          } else if offset == endIndex {
-            // We've found the node we want
-            return UnsafePath(node: node, slot: childSlot, childSlots: childSlots, offset: offset)
+          if handle.isLeaf {
+            offset += slot
+            targetSlot = slot
+            return .passUnretained(node.storage)
           } else {
-            startIndex = endIndex + 1
+            for i in 0...slot {
+              offset += handle[childAt: i].read({ $0.subtreeCount })
+            }
+            
+            let currentOffset = offset
+            let currentDepth = childSlots.depth
+            childSlots.append(UInt16(slot + 1))
+            
+            if let foundLater = search(in: handle[childAt: slot + 1]) {
+              return foundLater
+            } else {
+              childSlots.depth = currentDepth
+              targetSlot = slot
+              offset = currentOffset
+              
+              return .passUnretained(node.storage)
+            }
           }
+        } else {
+          // Start index exceeds node and is therefore not in this.
+          return nil
         }
-        
-        // TODO: convert into debug-only preconditionFaliure
-        preconditionFailure("In-bounds index not found within tree.")
-      }
-      
-      if let internalPath = internalPath { return Index(internalPath, forTree: self) }
+      })
     }
 
-    let path: UnsafePath = UnsafePath(
-      node: node,
-      slot: offset - startIndex,
-      childSlots: childSlots,
-      offset: offset
-    )
-    
-    return Index(path, forTree: self)
+    if let targetChild = search(in: self.root) {
+      return Index(UnsafePath(
+        node: targetChild,
+        slot: targetSlot,
+        childSlots: childSlots,
+        offset: offset
+      ), forTree: self)
+    } else {
+      return Index(nil, forTree: self)
+    }
   }
-
+  
 }

Sources/SortedCollections/BTree/_BTree.Index.swift

@@ -12,6 +12,8 @@
 extension _BTree {
 
   /// An index to an element of the BTree represented as a path.
+  /// - Warning: This has the capability to perform safety checks, however they must be explicitly be
+  ///     performed using the validation methods.
   @usableFromInline
   internal struct Index: Comparable {
     /// The path to the element in the BTree. A `nil` value indicates

Sources/SortedCollections/BTree/_BTree.UnsafePath.swift

@@ -13,11 +13,11 @@
 // TODO: potentially make operations mutating.
 
 extension _BTree {
-  /// Represents a specific element in a BTree. This holds strong references to the
+  /// Represents a specific element in a BTree. This does not hold any references to the
   /// element it points to.
   /// - Warning: Operations on this path will trap if the underlying node is deallocated.
   ///   and they become invalid if the tree is mutated, however this is not checked. For
-  ///   safety, use ``_BTree.Index`` instead
+  ///   safety, use ``_BTree.Index`` instead with its validation methods.
   @usableFromInline
   internal struct UnsafePath {
     @usableFromInline
@@ -68,7 +68,7 @@ extension _BTree {
       childSlots: Offsets,
       offset: Int
     ) {
-      self.init(node: node.storage, slot: slot, childSlots: childSlots, offset: offset)
+      self.init(node: .passUnretained(node.storage), slot: slot, childSlots: childSlots, offset: offset)
     }
 
     /// Creates a path representing a sequence of nodes to an element.
@@ -79,12 +79,12 @@ extension _BTree {
     ///   - index: The absolute offset of this path's element in the tree.
     @inlinable
     internal init(
-      node: Node.Storage,
+      node: Unmanaged<Node.Storage>,
       slot: Int,
       childSlots: Offsets,
       offset: Int
     ) {
-      self.node = .passUnretained(node)
+      self.node = node
       self.slot = slot
       self.childSlots = childSlots
       self.offset = offset
@@ -96,7 +96,7 @@ extension _BTree {
     @inlinable
     @inline(__always)
     internal var element: Element {
-      return self.readNode { $0[elementAt: self.slot] }
+      self.readNode { $0[elementAt: self.slot] }
     }
 
     /// Operators on a handle of the node
@@ -125,23 +125,9 @@ extension _BTree.UnsafePath: Equatable {
 // MARK: Comparable
 extension _BTree.UnsafePath: Comparable {
   /// Returns true if the first path points to an element before the second path
-  /// - Complexity: O(`log n`)
+  /// - Complexity: O(1)
   @inlinable
   public static func <(lhs: _BTree.UnsafePath, rhs: _BTree.UnsafePath) -> Bool {
-    for i in 0..<min(lhs.childSlots.depth, rhs.childSlots.depth) {
-      if lhs.childSlots[i] < rhs.childSlots[i] {
-        return true
-      }
-    }
-    
-    if lhs.childSlots.depth < rhs.childSlots.depth {
-      let rhsOffset = rhs.childSlots[lhs.childSlots.depth - 1]
-      return lhs.slot < rhsOffset
-    } else if rhs.childSlots.depth < lhs.childSlots.depth {
-      let lhsOffset = lhs.childSlots[rhs.childSlots.depth - 1]
-      return lhsOffset <= rhs.slot
-    } else {
-      return lhs.slot < rhs.slot
-    }
+    return lhs.offset < rhs.offset
   }
 }