SortedSet remove at index does not work

[pnewell]

Removing items from a SortedSet by index does remove the item, but does not update the count correctly.

Information

• Package version: main
• Platform version: macOS 15.3
• Swift version: 6.0.3

Checklist

• If possible, I've reproduced the issue using the main branch of this package.
• I've searched for existing GitHub issues.

Steps to Reproduce

let set = Set(0 ..< 40)
var sorted = SortedSet(set)
print(sorted.count)
print(sorted.remove(at: sorted.startIndex))
print(sorted)
print(sorted.count)

Expected behavior

40
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39]
0
39

Actual behavior

40
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39]
0
40

[pnewell]

Been playing around with this more, it ultimately leads to a runtime error the next time it checks its count. An easy way to see that is to convert to an array after the remove.

[natikgadzhi]

Looking at the code:

  // In SortedSet + Partial RangeReplaceableCollection.swift
  @inlinable
  @inline(__always)
  public mutating func remove(at index: Index) -> Element {
    index._index.ensureValid(forTree: self._root)
    return self._root.remove(at: index._index).key
  }

  // In SortedSet+BidirectionalCollection.swift
  @inlinable
  @inline(__always)
  public var count: Int { self._root.count }

I haven't ran the code locally yet, but looks like it passes both the count and the removal to the underlying BTree.

I'd start with making the failing test cases for both the SortedSet and the _BTree and see at what level the problem pops up.

BTree does have deletion tests though:

  func test_singleDeletion() {
    withEvery("size", in: [1, 2, 4, 8, 16, 32, 64, 128]) { size in
      withEvery("key", in: 0..<size) { key in
        btreeOfSize(size) { tree, kvs in
          tree.removeAnyElement(forKey: key)
          
          var comparisonKeys = Array(0..<size)
          comparisonKeys.remove(at: key)
          
          expectEqual(tree.count, size - 1)
          expectEqualElements(tree.map { $0.key }, comparisonKeys)
        }
      }
    }
  }