Implement initial SortedDictionary API and implementations for SortedSet.

Key Changes:
 - SortedSet
   - Implement BidirectionalCollection
 - SortedDictionary
   - Implement Codable
   - Implement Partial RangeReplaceableCollection
     - Includes remove*() operations
   - init(grouping:by:)
   - var keys: Keys (aka SortedSet<Key>)
   - func (compact)mapValues(_:)

Minor Changes:
 - Correct time documented complexities of various methods
 - Implement more variations of removal methods.

Author: vihanb

Sources/SortedCollections/BTree/_BTree.UnsafePath.swift

@@ -111,7 +111,8 @@ extension _BTree {
       node: Node,
       slot: Int,
       childSlots: Array<Int>,
-      offset: Int) {
+      offset: Int
+    ) {
       self.init(node: node.storage, slot: slot, childSlots: childSlots, offset: offset)
     }
 

Sources/SortedCollections/BTree/_BTree.swift

@@ -164,6 +164,22 @@ extension _BTree {
     return nil
   }
 
+  /// Verifies if the tree is balanced post-removal
+  /// - Warning: This does not invalidate indices.
+  @inlinable
+  @inline(__always)
+  internal mutating func _balanceRoot() {
+    if self.root.read({ $0.numElements == 0 && !$0.isLeaf }) {
+      let newRoot: Node = self.root.update { handle in
+        let newRoot = handle.moveChild(at: 0)
+        handle.drop()
+        return newRoot
+      }
+      
+      self.root = newRoot
+    }
+  }
+  
   /// Removes the key-value pair 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
@@ -179,22 +195,82 @@ extension _BTree {
     invalidateIndices()
 
     // TODO: Don't create CoW copy until needed.
-    // TODO: Handle root deletion
     let removedElement = self.root.update { $0.removeAny(key: key) }
 
     // Check if the tree height needs to be reduced
-    if self.root.read({ $0.numElements == 0 && !$0.isLeaf }) {
-      let newRoot: Node = self.root.update { handle in
-        let newRoot = handle.moveChild(at: 0)
-        handle.drop()
-        return newRoot
-      }
-      
-      self.root = newRoot
-    }
+    self._balanceRoot()
+    
+    return removedElement
+  }
+  
+  /// Removes the first element of a tree, if it exists.
+  ///
+  /// - Returns: The moved first element of the tree.
+  @inlinable
+  @discardableResult
+  internal mutating func popLast() -> Element? {
+    invalidateIndices()
 
+    if self.count == 0 { return nil }
+    
+    let removedElement = self.root.update { $0.popLastElement() }
+    self._balanceRoot()
+    return removedElement
+  }
+  
+  /// Removes the first element of a tree, if it exists.
+  ///
+  /// - Returns: The moved first element of the tree.
+  @inlinable
+  @discardableResult
+  internal mutating func popFirst() -> Element? {
+    invalidateIndices()
+    
+    if self.count == 0 { return nil }
+    
+    let removedElement = self.root.update { $0.popFirstElement() }
+    self._balanceRoot()
+    return removedElement
+  }
+  
+  /// Removes the element of a tree at a given offset.
+  ///
+  /// - Parameter offset: the offset which must be in-bounds.
+  /// - Returns: The moved element of the tree
+  @inlinable
+  @discardableResult
+  internal mutating func remove(at offset: Int) -> Element {
+    invalidateIndices()
+    let removedElement = self.root.update { $0.remove(at: offset) }
+    self._balanceRoot()
     return removedElement
   }
+  
+  /// 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
+  @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>) {
+    guard let startPath = bounds.lowerBound.path else { preconditionFailure("Index out of bounds.") }
+    guard let _ = bounds.upperBound.path else { preconditionFailure("Index out of bounds.") }
+    
+    let rangeSize = self.distance(from: bounds.lowerBound, to: bounds.upperBound)
+    let startOffset = startPath.offset
+    
+    for _ in 0..<rangeSize {
+      self.remove(at: startOffset)
+    }
+  }
 }
 
 // MARK: Read Operations
@@ -234,7 +310,6 @@ extension _BTree {
     return false
   }
 
-  
   /// Returns the value corresponding to the first found instance of the key.
   ///
   /// This may not be the first instance of the key. This is marginally more efficient

Sources/SortedCollections/BTree/_Node.UnsafeHandle+Deletion.swift

@@ -36,8 +36,8 @@ extension _Node.UnsafeHandle {
       } else {
         // Deletion within an internal node
 
-        // TODO: potentially be smarter about using the predecessor or successor.0
-        let predecessor = self[childAt: slot].update { $0.popElement() }
+        // TODO: potentially be smarter about using the predecessor or successor.
+        let predecessor = self[childAt: slot].update { $0.popLastElement() }
 
         // Reduce the element count.
         self.numTotalElements -= 1
@@ -84,6 +84,77 @@ extension _Node.UnsafeHandle {
     }
   }
 
+  /// Removes the element of the tree rooted at this node, at a given offset.
+  ///
+  /// This may leave the node it is called upon unbalanced so it is important to
+  /// ensure the tree above this is balanced. This does adjust child counts
+  ///
+  /// - Parameter offset: the offset which must be in-bounds.
+  /// - Returns: The moved element of the tree
+  @inlinable
+  @inline(__always)
+  internal func remove(at offset: Int) -> _Node.Element {
+    assertMutable()
+    assert(0 <= offset && offset < self.numTotalElements, "Cannot remove with out-of-bounds offset.")
+    
+    if self.isLeaf {
+      return self.removeElement(at: offset)
+    } else {
+      var startIndex = 0
+      for childSlot in 0..<self.numChildren {
+        let endIndex = startIndex + self[childAt: childSlot].read { $0.numTotalElements }
+        
+        if offset < endIndex {
+          return self[childAt: childSlot].update { $0.remove(at: offset - startIndex) }
+        } else if offset == endIndex {
+          let predecessor = self[childAt: childSlot].update { $0.popLastElement() }
+          
+          self.numTotalElements -= 1
+          
+          // Replace the current element with the predecessor.
+          let element = self.moveElement(at: childSlot)
+          self.setElement(predecessor, at: childSlot)
+          
+          // Balance the predecessor child slot, as the pop operation may have
+          // brought it out of balance.
+          self.balance(at: childSlot)
+          
+          return element
+        } else {
+          startIndex = endIndex + 1
+        }
+      }
+      
+      preconditionFailure("B-Tree in invalid state.")
+    }
+  }
+  
+  /// Removes the first element of a tree, balancing the tree.
+  ///
+  /// This may leave the node it is called upon unbalanced so it is important to
+  /// ensure the tree above this is balanced. This does adjust child counts
+  ///
+  /// - Returns: The moved first element of the tree.
+  @inlinable
+  @inline(__always)
+  internal func popFirstElement() -> _Node.Element {
+    assertMutable()
+    
+    if self.isLeaf {
+      // At a leaf, it is trivial to pop the last element
+      // removeElement(at:) automatically updates the counts.
+      return self.removeElement(at: 0)
+    } else {
+      // Remove the subtree's element
+      let poppedElement = self[childAt: 0].update { $0.popFirstElement() }
+      
+      self.numTotalElements -= 1
+      
+      self.balance(at: self.numChildren - 1)
+      return poppedElement
+    }
+  }
+  
   /// Removes the last element of a tree, balancing the tree.
   ///
   /// This may leave the node it is called upon unbalanced so it is important to
@@ -92,16 +163,16 @@ extension _Node.UnsafeHandle {
   /// - Returns: The moved last element of the tree.
   @inlinable
   @inline(__always)
-  internal func popElement() -> _Node.Element {
+  internal func popLastElement() -> _Node.Element {
     assertMutable()
 
     if self.isLeaf {
       // At a leaf, it is trivial to pop the last element
-      // removeElement(at:) automatically updates the counts.
+      // popLastElement(at:) automatically updates the counts.
       return self.removeElement(at: self.numElements - 1)
     } else {
       // Remove the subtree's element
-      let poppedElement = self[childAt: self.numChildren - 1].update { $0.popElement() }
+      let poppedElement = self[childAt: self.numChildren - 1].update { $0.popLastElement() }
 
       self.numTotalElements -= 1
 

Sources/SortedCollections/SortedDictionary/SortedDictionary+BidirectionalCollection.swift

@@ -80,10 +80,4 @@ extension SortedDictionary: BidirectionalCollection {
     i._index.ensureValid(for: self._root)
     return Index(self._root.index(i._index, offsetBy: distance))
   }
-  
-  @inlinable
-  public subscript(position: Index) -> Element {
-    position._index.ensureValid(for: self._root)
-    return self._root[position._index]
-  }
 }

Sources/SortedCollections/SortedDictionary/SortedDictionary+Codable.swift

@@ -0,0 +1,69 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift Collections open source project
+//
+// Copyright (c) 2021 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
+//
+//===----------------------------------------------------------------------===//
+
+extension SortedDictionary: Encodable where Key: Codable, Value: Codable {
+  /// Encodes the contents of this dictionary into the given encoder.
+  ///
+  /// The dictionary's contents are encoded as alternating key-value pairs in
+  /// an unkeyed container.
+  ///
+  /// This function throws an error if any values are invalid for the given
+  /// encoder's format.
+  ///
+  /// - Note: Unlike the standard `Dictionary` type, sorted dictionaries
+  ///    always encode themselves into an unkeyed container, because
+  ///    `Codable`'s keyed containers do not guarantee that they preserve the
+  ///    ordering of the items they contain. (And in popular encoding formats,
+  ///    keyed containers tend to map to unordered data structures -- e.g.,
+  ///    JSON's "object" construct is explicitly unordered.)
+  ///
+  /// - Parameter encoder: The encoder to write data to.
+  @inlinable
+  public func encode(to encoder: Encoder) throws {
+    // Encode contents as an array of alternating key-value pairs.
+    var container = encoder.unkeyedContainer()
+    for (key, value) in self {
+      try container.encode(key)
+      try container.encode(value)
+    }
+  }
+}
+
+extension SortedDictionary: Decodable where Key: Decodable, Value: Decodable {
+  @inlinable
+  public init(from decoder: Decoder) throws {
+    // We expect to be encoded as an array of alternating key-value pairs.
+    var container = try decoder.unkeyedContainer()
+
+    self.init()
+    while !container.isAtEnd {
+      let key = try container.decode(Key.self)
+      
+      guard !container.isAtEnd else {
+        throw DecodingError.dataCorrupted(
+          DecodingError.Context(
+            codingPath: container.codingPath,
+            debugDescription: "Unkeyed container reached end before value in key-value pair"
+          )
+        )
+      }
+      
+      let value = try container.decode(Value.self)
+      let oldValue = self.updateValue(value, forKey: key)
+      if oldValue != nil {
+        let context = DecodingError.Context(
+          codingPath: container.codingPath,
+          debugDescription: "Duplicate key at offset \(container.currentIndex - 1)")
+        throw DecodingError.dataCorrupted(context)
+      }
+    }
+  }
+}

Sources/SortedCollections/SortedDictionary/SortedDictionary+Equatable.swift

@@ -20,7 +20,7 @@ extension SortedDictionary: Equatable where Value: Equatable {
   /// - Parameters:
   ///   - lhs: A value to compare.
   ///   - rhs: Another value to compare.
-  /// - Complexity: O(`n log n`)
+  /// - Complexity: O(`n`)
   @inlinable
   public static func ==(lhs: SortedDictionary<Key, Value>, rhs: SortedDictionary<Key, Value>) -> Bool {
     if lhs.count != rhs.count { return false }

Sources/SortedCollections/SortedDictionary/SortedDictionary+Hashable.swift

@@ -14,7 +14,7 @@ extension SortedDictionary: Hashable where Key: Hashable, Value: Hashable {
   /// into the givenhasher.
   /// - Parameter hasher: The hasher to use when combining
   ///     the components of this instance.
-  /// - Complexity: O(`n log n`)
+  /// - Complexity: O(`n`)
   @inlinable
   public func hash(into hasher: inout Hasher) {
     hasher.combine(self.count)