diff --git a/.swiftpm/xcode/xcshareddata/xcschemes/SortedCollections.xcscheme b/.swiftpm/xcode/xcshareddata/xcschemes/SortedCollections.xcscheme new file mode 100644 index 000000000..a660e5897 --- /dev/null +++ b/.swiftpm/xcode/xcshareddata/xcschemes/SortedCollections.xcscheme @@ -0,0 +1,87 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/.swiftpm/xcode/xcshareddata/xcschemes/swift-collections-Package.xcscheme b/.swiftpm/xcode/xcshareddata/xcschemes/swift-collections-Package.xcscheme index bf3d7efa4..a3af5c698 100644 --- a/.swiftpm/xcode/xcshareddata/xcschemes/swift-collections-Package.xcscheme +++ b/.swiftpm/xcode/xcshareddata/xcschemes/swift-collections-Package.xcscheme @@ -37,14 +37,14 @@ @@ -56,9 +56,9 @@ buildForAnalyzing = "YES"> @@ -70,9 +70,9 @@ buildForAnalyzing = "YES"> @@ -106,29 +106,29 @@ + buildForRunning = "YES" + buildForProfiling = "YES" + buildForArchiving = "YES" + buildForAnalyzing = "YES"> + buildForRunning = "YES" + buildForProfiling = "YES" + buildForArchiving = "YES" + buildForAnalyzing = "YES"> @@ -140,9 +140,9 @@ buildForAnalyzing = "YES"> @@ -194,6 +194,16 @@ ReferencedContainer = "container:"> + + + + + + + + + + + + insert", + input: [Int].self + ) { input in + input.withUnsafeBufferPointer { buffer in + cpp_map_insert_integers(buffer.baseAddress, buffer.count) + } + } + + self.add( + title: "std::map successful find", + input: ([Int], [Int]).self + ) { input, lookups in + let map = CppMap(input) + return { timer in + lookups.withUnsafeBufferPointer { buffer in + cpp_map_lookups(map.ptr, buffer.baseAddress, buffer.count) + } + } + } } } diff --git a/Benchmarks/Benchmarks/SortedDictionaryBenchmarks.swift b/Benchmarks/Benchmarks/SortedDictionaryBenchmarks.swift new file mode 100644 index 000000000..5af4c381d --- /dev/null +++ b/Benchmarks/Benchmarks/SortedDictionaryBenchmarks.swift @@ -0,0 +1,303 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +import CollectionsBenchmark +import SortedCollections + +extension Benchmark { + public mutating func addSortedDictionaryBenchmarks() { + self.add( + title: "SortedDictionary init(keysWithValues:)", + input: [Int].self + ) { input in + let keysAndValues = input.map { (key: $0, value: 2 * $0) } + + return { timer in + blackHole(SortedDictionary(keysWithValues: keysAndValues)) + } + } + + self.add( + title: "SortedDictionary init(sortedKeysWithValues:)", + input: Int.self + ) { input in + let keysAndValues = (0.. sort, then init(sortedKeysWithValues:)", + input: [Int].self + ) { input in + return { timer in + var keysAndValues = input.map { (key: $0, value: 2 * $0) } + + timer.measure { + keysAndValues.sort(by: { $0.key < $1.key }) + blackHole(SortedDictionary(sortedKeysWithValues: keysAndValues)) + } + } + } + + self.add( + title: "SortedDictionary sequential iteration", + input: [Int].self + ) { input in + let keysAndValues = input.lazy.map { (key: $0, value: 2 * $0) } + let d = SortedDictionary(keysWithValues: keysAndValues) + + return { timer in + for item in d { + blackHole(item) + } + } + } + + self.add( + title: "SortedDictionary index-based iteration", + input: [Int].self + ) { input in + let keysAndValues = input.lazy.map { (key: $0, value: 2 * $0) } + let d = SortedDictionary(keysWithValues: keysAndValues) + + return { timer in + var i = d.startIndex + while i != d.endIndex { + blackHole(d[i]) + d.formIndex(after: &i) + } + } + } + + self.add( + title: "SortedDictionary offset-based iteration", + input: [Int].self + ) { input in + let keysAndValues = input.lazy.map { (key: $0, value: 2 * $0) } + let d = SortedDictionary(keysWithValues: keysAndValues) + + return { timer in + for offset in 0.. forEach iteration", + input: [Int].self + ) { input in + let keysAndValues = input.lazy.map { (key: $0, value: 2 * $0) } + let d = SortedDictionary(keysWithValues: keysAndValues) + + return { timer in + d.forEach({ blackHole($0) }) + } + } + + self.add( + title: "SortedDictionary.Keys sequential iteration", + input: [Int].self + ) { input in + let keysAndValues = input.lazy.map { (key: $0, value: 2 * $0) } + let d = SortedDictionary(keysWithValues: keysAndValues) + + return { timer in + for item in d.keys { + blackHole(item) + } + } + } + + self.add( + title: "SortedDictionary.Values sequential iteration", + input: [Int].self + ) { input in + let keysAndValues = input.lazy.map { (key: $0, value: 2 * $0) } + let d = SortedDictionary(keysWithValues: keysAndValues) + + return { timer in + for item in d.values { + blackHole(item) + } + } + } + + self.add( + title: "SortedDictionary subscript, successful lookups", + input: ([Int], [Int]).self + ) { input, lookups in + let sortedDictionary = SortedDictionary( + keysWithValues: input.map { ($0, 2 * $0) }) + + return { timer in + for key in lookups { + precondition(sortedDictionary[key] == key * 2) + } + } + } + + self.add( + title: "SortedDictionary subscript, unsuccessful lookups", + input: ([Int], [Int]).self + ) { input, lookups in + let sortedDictionary = SortedDictionary( + keysWithValues: input.map { ($0, 2 * $0) }) + + let c = input.count + return { timer in + for key in lookups { + precondition(sortedDictionary[key + c] == nil) + } + } + } + + self.add( + title: "SortedDictionary subscript, setter append", + input: [Int].self + ) { input in + let keysAndValues = input.lazy.map { (key: $0, value: 2 * $0) } + var sortedDictionary = SortedDictionary() + + return { timer in + for (key, value) in keysAndValues { + sortedDictionary[key] = value + } + blackHole(sortedDictionary) + } + } + + self.add( + title: "SortedDictionary subscript, setter noop", + input: ([Int], [Int]).self + ) { input, lookups in + return { timer in + var d = SortedDictionary( + keysWithValues: input.map { ($0, 2 * $0) }) + + let c = input.count + timer.measure { + for i in lookups { + d[i + c] = nil + } + } + precondition(d.count == input.count) + blackHole(d) + } + } + + self.add( + title: "SortedDictionary subscript, setter update", + input: ([Int], [Int]).self + ) { input, lookups in + return { timer in + var d = SortedDictionary( + keysWithValues: input.map { ($0, 2 * $0) }) + + timer.measure { + for i in lookups { + d[i] = 0 + } + } + precondition(d.count == input.count) + blackHole(d) + } + } + + self.add( + title: "SortedDictionary subscript, setter remove", + input: ([Int], [Int]).self + ) { input, lookups in + return { timer in + var d = SortedDictionary( + keysWithValues: input.map { ($0, 2 * $0) }) + + timer.measure { + for i in lookups { + d[i] = nil + } + } + precondition(d.count == 0) + blackHole(d) + } + } + + self.add( + title: "SortedDictionary subscript, _modify insert", + input: ([Int], [Int]).self + ) { input, lookups in + return { timer in + var d = SortedDictionary() + + @inline(__always) + func modify(_ i: inout Int?, to value: Int?) { + i = value + } + + timer.measure { + for i in lookups { + modify(&d[i], to: i * 2) + } + } + + precondition(d.count == input.count) + blackHole(d) + } + } + + self.add( + title: "SortedDictionary subscript, _modify update", + input: ([Int], [Int]).self + ) { input, lookups in + return { timer in + var d = SortedDictionary( + keysWithValues: input.map { ($0, 2 * $0) }) + + timer.measure { + for i in lookups { + d[i]! *= 2 + } + } + precondition(d.count == input.count) + blackHole(d) + } + } + + self.add( + title: "SortedDictionary subscript, _modify remove", + input: ([Int], [Int]).self + ) { input, lookups in + return { timer in + var d = SortedDictionary( + keysWithValues: input.map { ($0, 2 * $0) }) + + @inline(__always) + func modify(_ i: inout Int?, to value: Int?) { + i = value + } + + timer.measure { + for i in lookups { + modify(&d[i], to: nil) + } + } + + precondition(d.count == 0) + blackHole(d) + } + } + } +} diff --git a/Benchmarks/Benchmarks/SortedSetBenchmarks.swift b/Benchmarks/Benchmarks/SortedSetBenchmarks.swift new file mode 100644 index 000000000..be4555d4d --- /dev/null +++ b/Benchmarks/Benchmarks/SortedSetBenchmarks.swift @@ -0,0 +1,207 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +import CollectionsBenchmark +import SortedCollections + +extension Benchmark { + public mutating func addSortedSetBenchmarks() { + self.addSimple( + title: "SortedSet init from range", + input: Int.self + ) { size in + blackHole(SortedSet(0 ..< size)) + } + + self.addSimple( + title: "SortedSet init(sortedElements:) from range", + input: Int.self + ) { size in + blackHole(SortedSet(sortedElements: 0 ..< size)) + } + self.add( + title: "SortedSet sequential iteration", + input: [Int].self + ) { input in + let set = SortedSet(input) + return { timer in + for i in set { + blackHole(i) + } + } + } + + self.add( + title: "SortedSet forEach iteration", + input: [Int].self + ) { input in + let set = SortedSet(input) + return { timer in + set.forEach { i in + blackHole(i) + } + } + } + + self.add( + title: "SortedSet successful contains", + input: ([Int], [Int]).self + ) { input, lookups in + let set = SortedSet(input) + return { timer in + for i in lookups { + precondition(set.contains(i)) + } + } + } + + self.add( + title: "SortedSet unsuccessful contains", + input: ([Int], [Int]).self + ) { input, lookups in + let set = SortedSet(input) + let lookups = lookups.map { $0 + input.count } + return { timer in + for i in lookups { + precondition(!set.contains(i)) + } + } + } + + self.addSimple( + title: "SortedSet insertions", + input: [Int].self + ) { input in + var set: SortedSet = [] + for i in input { + set.insert(i) + } + precondition(set.count == input.count) + blackHole(set) + } + + self.add( + title: "SortedSet remove", + input: ([Int], [Int]).self + ) { input, removals in + return { timer in + var set = SortedSet(input) + timer.measure { + for i in removals { + set.remove(i) + } + } + precondition(set.isEmpty) + blackHole(set) + } + } + + self.add( + title: "SortedSet removeLast", + input: Int.self + ) { size in + return { timer in + var set = SortedSet(0 ..< size) + timer.measure { + for _ in 0 ..< size { + set.removeLast() + } + } + precondition(set.isEmpty) + blackHole(set) + } + } + + self.add( + title: "SortedSet removeFirst", + input: Int.self + ) { size in + return { timer in + var set = SortedSet(0 ..< size) + timer.measure { + for _ in 0 ..< size { + set.removeFirst() + } + } + precondition(set.isEmpty) + blackHole(set) + } + } + + let overlaps: [(String, (Int) -> Int)] = [ + ("0%", { c in c }), + ("25%", { c in 3 * c / 4 }), + ("50%", { c in c / 2 }), + ("75%", { c in c / 4 }), + ("100%", { c in 0 }), + ] + + // SetAlgebra operations with Self + do { + for (percentage, start) in overlaps { + self.add( + title: "SortedSet union with Self (\(percentage) overlap)", + input: [Int].self + ) { input in + let start = start(input.count) + let a = SortedSet(input) + let b = SortedSet(start ..< start + input.count) + return { timer in + blackHole(a.union(b)) + } + } + } + + for (percentage, start) in overlaps { + self.add( + title: "SortedSet intersection with Self (\(percentage) overlap)", + input: [Int].self + ) { input in + let start = start(input.count) + let a = SortedSet(input) + let b = SortedSet(start ..< start + input.count) + return { timer in + blackHole(a.intersection(b)) + } + } + } + + for (percentage, start) in overlaps { + self.add( + title: "SortedSet symmetricDifference with Self (\(percentage) overlap)", + input: [Int].self + ) { input in + let start = start(input.count) + let a = SortedSet(input) + let b = SortedSet(start ..< start + input.count) + return { timer in + blackHole(a.symmetricDifference(b)) + } + } + } + + for (percentage, start) in overlaps { + self.add( + title: "SortedSet subtracting Self (\(percentage) overlap)", + input: [Int].self + ) { input in + let start = start(input.count) + let a = SortedSet(input) + let b = SortedSet(start ..< start + input.count) + return { timer in + blackHole(a.subtracting(b)) + } + } + } + } + + } +} diff --git a/Benchmarks/CppBenchmarks/include/MapBenchmarks.h b/Benchmarks/CppBenchmarks/include/MapBenchmarks.h new file mode 100644 index 000000000..f02db76be --- /dev/null +++ b/Benchmarks/CppBenchmarks/include/MapBenchmarks.h @@ -0,0 +1,40 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#ifndef CPPBENCHMARKS_MAP_BENCHMARKS_H +#define CPPBENCHMARKS_MAP_BENCHMARKS_H + +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/// Create a std::map, populating it with data from the supplied buffer. +/// Returns an opaque pointer to the created instance. +extern void *cpp_map_create(const intptr_t *start, size_t count); + +/// Destroys an ordered map previously returned by `cpp_map_create`. +extern void cpp_map_destroy(void *ptr); + +extern void cpp_map_insert_integers(const intptr_t *start, size_t count); + +extern void cpp_map_lookups(void *ptr, const intptr_t *start, size_t count); +extern void cpp_map_subscript(void *ptr, const intptr_t *start, size_t count); + +#ifdef __cplusplus +} +#endif + + +#endif /* CPPBENCHMARKS_MAP_BENCHMARKS_H */ diff --git a/Benchmarks/CppBenchmarks/include/module.modulemap b/Benchmarks/CppBenchmarks/include/module.modulemap index 7fbcf9332..900b0f1ab 100644 --- a/Benchmarks/CppBenchmarks/include/module.modulemap +++ b/Benchmarks/CppBenchmarks/include/module.modulemap @@ -5,6 +5,7 @@ module CppBenchmarks { header "DequeBenchmarks.h" header "UnorderedSetBenchmarks.h" header "UnorderedMapBenchmarks.h" + header "MapBenchmarks.h" header "PriorityQueueBenchmarks.h" export * } diff --git a/Benchmarks/CppBenchmarks/src/MapBenchmarks.cpp b/Benchmarks/CppBenchmarks/src/MapBenchmarks.cpp new file mode 100644 index 000000000..60488b78c --- /dev/null +++ b/Benchmarks/CppBenchmarks/src/MapBenchmarks.cpp @@ -0,0 +1,70 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include "MapBenchmarks.h" +#include +#include +#include "utils.h" + +typedef std::map custom_map; + +void * +cpp_map_create(const intptr_t *start, size_t count) +{ + auto map = new custom_map(); + for (size_t i = 0; i < count; ++i) { + map->insert({start[i], 2 * start[i]}); + } + return map; +} + +void +cpp_map_destroy(void *ptr) +{ + delete static_cast(ptr); +} + +void +cpp_map_insert_integers(const intptr_t *start, size_t count) +{ + auto map = custom_map(); + auto end = start + count; + for (auto p = start; p != end; ++p) { + auto v = *identity(p); + map.insert({ v, 2 * v }); + } + black_hole(&map); +} + +__attribute__((noinline)) +auto find(custom_map* map, intptr_t value) +{ + return map->find(value); +} + +void +cpp_map_lookups(void *ptr, const intptr_t *start, size_t count) +{ + auto map = static_cast(ptr); + for (auto it = start; it < start + count; ++it) { + auto isCorrect = find(map, *it)->second == *it * 2; + if (!isCorrect) { abort(); } + } +} + +void +cpp_map_subscript(void *ptr, const intptr_t *start, size_t count) +{ + auto map = static_cast(ptr); + for (auto it = start; it < start + count; ++it) { + black_hole((*map)[*it]); + } +} diff --git a/Benchmarks/swift-collections-benchmark/main.swift b/Benchmarks/swift-collections-benchmark/main.swift index d2096d384..103ad3ca5 100644 --- a/Benchmarks/swift-collections-benchmark/main.swift +++ b/Benchmarks/swift-collections-benchmark/main.swift @@ -19,6 +19,8 @@ benchmark.addDictionaryBenchmarks() benchmark.addDequeBenchmarks() benchmark.addOrderedSetBenchmarks() benchmark.addOrderedDictionaryBenchmarks() +benchmark.addSortedSetBenchmarks() +benchmark.addSortedDictionaryBenchmarks() benchmark.addHeapBenchmarks() benchmark.addCppBenchmarks() #if os(macOS) || os(iOS) || os(watchOS) || os(tvOS) diff --git a/Package.swift b/Package.swift index ace16942c..3ab9cdc80 100644 --- a/Package.swift +++ b/Package.swift @@ -54,6 +54,7 @@ let package = Package( .library(name: "Collections", targets: ["Collections"]), .library(name: "DequeModule", targets: ["DequeModule"]), .library(name: "OrderedCollections", targets: ["OrderedCollections"]), + .library(name: "SortedCollections", targets: ["SortedCollections"]), .library(name: "PriorityQueueModule", targets: ["PriorityQueueModule"]), ], dependencies: [ @@ -66,6 +67,7 @@ let package = Package( dependencies: [ "DequeModule", "OrderedCollections", + "SortedCollections", "PriorityQueueModule", ], path: "Sources/Collections", @@ -132,7 +134,16 @@ let package = Package( name: "OrderedCollectionsTests", dependencies: ["OrderedCollections", "CollectionsTestSupport"], swiftSettings: settings), - + + // SortedDictionary + .target( + name: "SortedCollections", + swiftSettings: settings), + .testTarget( + name: "SortedCollectionsTests", + dependencies: ["SortedCollections", "CollectionsTestSupport"], + swiftSettings: settings), + // PriorityQueue .target( name: "PriorityQueueModule", diff --git a/Sources/SortedCollections/BTree/_BTree+BidirectionalCollection.swift b/Sources/SortedCollections/BTree/_BTree+BidirectionalCollection.swift new file mode 100644 index 000000000..f676bafaa --- /dev/null +++ b/Sources/SortedCollections/BTree/_BTree+BidirectionalCollection.swift @@ -0,0 +1,205 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +// This contains `_BTree`'s general implementation of BidirectionalCollection. +// These operations are bounds in contrast to most other methods on _BTree as +// they are designed to be easily propogated to a higher-level data type. +// However, they still do not perform index validation + +extension _BTree: BidirectionalCollection { + /// The total number of elements contained within the BTree + /// - Complexity: O(1) + @inlinable + @inline(__always) + internal var count: Int { self.root.storage.header.subtreeCount } + + /// A Boolean value that indicates whether the BTree is empty. + @inlinable + @inline(__always) + internal var isEmpty: Bool { self.count == 0 } + + // TODO: further consider O(1) implementation + /// Locates the first element and returns a proper path to it, or nil if the BTree is empty. + /// - Complexity: O(`log n`) + @inlinable + internal var startIndex: Index { + if count == 0 { return endIndex } + var depth: Int8 = 0 + var currentNode: Unmanaged = .passUnretained(self.root.storage) + while true { + let shouldStop: Bool = currentNode._withUnsafeGuaranteedRef { + $0.read { handle in + if handle.isLeaf { + return true + } else { + depth += 1 + currentNode = .passUnretained(handle[childAt: 0].storage) + return false + } + } + } + + if shouldStop { break } + } + + return Index( + node: currentNode, + slot: 0, + childSlots: _FixedSizeArray(repeating: 0, depth: depth), + offset: 0, + forTree: self + ) + } + + /// Returns a sentinel value for the last element + /// - Complexity: O(1) + @inlinable + internal var endIndex: Index { + Index( + node: .passUnretained(self.root.storage), + slot: -1, + childSlots: Index.Offsets(repeating: 0), + offset: self.count, + forTree: self + ) + } + + /// Returns the distance between two indices. + /// - 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. + /// - Complexity: O(1) + @inlinable + internal func distance(from start: Index, to end: Index) -> Int { + return end.offset - start.offset + } + + /// Replaces the given index with its successor. + /// - Parameter index: A valid index of the collection. i must be less than endIndex. + /// - Complexity: O(`log n`) in the worst-case. + @inlinable + internal func formIndex(after index: inout Index) { + precondition(index.offset < self.count, + "Attempt to advance out of collection bounds.") + + // TODO: this might be redundant given the fact the same (but generalized) + // logic is implemented in offsetBy + let shouldSeekWithinLeaf = index.readNode { + $0.isLeaf && _fastPath(index.slot + 1 < $0.elementCount) + } + + if shouldSeekWithinLeaf { + // Continue searching within the same leaf + index.slot += 1 + index.offset += 1 + } else { + self.formIndex(&index, offsetBy: 1) + } + } + + /// Returns the position immediately after the given index. + /// - Parameter i: A valid index of the collection. i must be less than endIndex. + /// - Returns: The index value immediately after i. + /// - Complexity: O(`log n`) in the worst-case. + @inlinable + internal func index(after i: Index) -> Index { + var newIndex = i + self.formIndex(after: &newIndex) + return newIndex + } + + /// Replaces the given index with its predecessor. + /// - Parameter index: A valid index of the collection. i must be greater than startIndex. + /// - Complexity: O(`log n`) in the worst-case. + @inlinable + internal func formIndex(before index: inout Index) { + precondition(!self.isEmpty && index.offset != 0, + "Attempt to advance out of collection bounds.") + self.formIndex(&index, offsetBy: -1) + } + + /// Returns the position immediately before the given index. + /// - Parameter i: A valid index of the collection. i must be greater than startIndex. + /// - Returns: The index value immediately before i. + /// - Complexity: O(`log n`) in the worst-case. + @inlinable + internal func index(before i: Index) -> Index { + var newIndex = i + self.formIndex(before: &newIndex) + return newIndex + } + + /// Offsets the given index by the specified distance. + /// + /// The value passed as distance must not offset i beyond the bounds of the collection. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - Complexity: O(`log n`) in the worst-case. + @inlinable + internal func formIndex(_ i: inout Index, offsetBy distance: Int) { + let newIndex = i.offset + distance + precondition(0 <= newIndex && newIndex <= self.count, + "Attempt to advance out of collection bounds.") + + if newIndex == self.count { + i = endIndex + return + } + + // TODO: optimization for searching within children + + if i != endIndex && i.readNode({ $0.isLeaf }) { + // Check if the target element will be in the same node + let targetSlot = i.slot + distance + if 0 <= targetSlot && targetSlot < i.readNode({ $0.elementCount }) { + i.slot = targetSlot + i.offset = newIndex + return + } + } + + // Otherwise, re-seek + i = self.index(atOffset: newIndex) + } + + /// Returns an index that is the specified distance from the given index. + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - Returns: An index offset by `distance` from the index `i`. If `distance` + /// is positive, this is the same value as the result of `distance` calls to + /// `index(after:)`. If `distance` is negative, this is the same value as the + /// result of `abs(distance)` calls to `index(before:)`. + @inlinable + internal func index(_ i: Index, offsetBy distance: Int) -> Index { + var newIndex = i + self.formIndex(&newIndex, offsetBy: distance) + return newIndex + } + + @inlinable + @inline(__always) + internal subscript(index: Index) -> Element { + // Ensure we don't attempt to dereference the endIndex + precondition(index != endIndex, "Attempt to subscript out of range index.") + return index.element + } + + @inlinable + @inline(__always) + internal subscript(bounds: Range) -> SubSequence { + return SubSequence(base: self, bounds: bounds) + } +} diff --git a/Sources/SortedCollections/BTree/_BTree+CustomDebugStringConvertible.swift b/Sources/SortedCollections/BTree/_BTree+CustomDebugStringConvertible.swift new file mode 100644 index 000000000..a0b0489f1 --- /dev/null +++ b/Sources/SortedCollections/BTree/_BTree+CustomDebugStringConvertible.swift @@ -0,0 +1,25 @@ +//===----------------------------------------------------------------------===// +// +// 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 _BTree: CustomDebugStringConvertible { + #if DEBUG + /// A textual representation of this instance, suitable for debugging. + public var debugDescription: String { + return "BTree<\(Key.self), \(Value.self)>\n" + + self.root.read { String(reflecting: $0) } + } + #else + /// A textual representation of this instance, suitable for debugging. + public var debugDescription: String { + return "BTree<\(Key.self), \(Value.self)>(\(self.root))" + } + #endif // DEBUG +} diff --git a/Sources/SortedCollections/BTree/_BTree+CustomReflectable.swift b/Sources/SortedCollections/BTree/_BTree+CustomReflectable.swift new file mode 100644 index 000000000..30721d572 --- /dev/null +++ b/Sources/SortedCollections/BTree/_BTree+CustomReflectable.swift @@ -0,0 +1,18 @@ +//===----------------------------------------------------------------------===// +// +// 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 _BTree: CustomReflectable { + /// The custom mirror for this instance. + @inlinable + internal var customMirror: Mirror { + Mirror(self, unlabeledChildren: self, displayStyle: .dictionary) + } +} diff --git a/Sources/SortedCollections/BTree/_BTree+Invariants.swift b/Sources/SortedCollections/BTree/_BTree+Invariants.swift new file mode 100644 index 000000000..b2684a0c9 --- /dev/null +++ b/Sources/SortedCollections/BTree/_BTree+Invariants.swift @@ -0,0 +1,102 @@ +//===----------------------------------------------------------------------===// +// +// 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 _BTree { + #if COLLECTIONS_INTERNAL_CHECKS + @inline(never) + @discardableResult + fileprivate func checkInvariants( + for node: Node, + expectedDepth: Int, + isRoot: Bool = false + ) -> ( + minimum: Key?, + maximum: Key? + ) { + node.read { handle in + assert(handle.depth == expectedDepth, "Node depth mismatch.") + assert(isRoot || handle.elementCount > 0, "Node cannot be empty") + + if handle.elementCount > 1 { + for i in 0..<(handle.elementCount - 1) { + assert(handle[keyAt: i] <= handle[keyAt: i + 1], + "Node keys out of order.") + } + } + + if handle.isLeaf { + assert(handle.elementCount == handle.subtreeCount, + "Element and subtree count should match for leaves.") + assert(handle.depth == 0, "Non-zero depth for leaf.") + assert(isRoot || handle.isBalanced, "Unbalanced node.") + + if handle.elementCount > 0 { + return ( + minimum: handle[keyAt: 0], + maximum: handle[keyAt: handle.elementCount - 1] + ) + } else { + return (nil, nil) + } + } else { + var totalCount = 0 + var subtreeMinimum: Key! + var subtreeMaximum: Key! + + for i in 0..= handle[keyAt: i - 1], + "Last subtree must be greater than or equal to last key.") + } else { + assert(maximum! <= handle[keyAt: i], + "Subtree must be less than or equal to corresponding key.") + } + + totalCount += handle[childAt: i].read { $0.subtreeCount } + } + + assert(handle.subtreeCount == handle.elementCount + totalCount, + "Subtree count mismatch.") + + return ( + minimum: subtreeMinimum, + maximum: subtreeMaximum + ) + } + } + } + + @inline(never) + @usableFromInline + internal func checkInvariants() { + checkInvariants( + for: root, + expectedDepth: root.storage.header.depth, + isRoot: true + ) + } + #else + @inlinable + @inline(__always) + internal func checkInvariants() {} + #endif // COLLECTIONS_INTERNAL_CHECKS +} diff --git a/Sources/SortedCollections/BTree/_BTree+Partial RangeReplaceableCollection.swift b/Sources/SortedCollections/BTree/_BTree+Partial RangeReplaceableCollection.swift new file mode 100644 index 000000000..57eada023 --- /dev/null +++ b/Sources/SortedCollections/BTree/_BTree+Partial RangeReplaceableCollection.swift @@ -0,0 +1,141 @@ +//===----------------------------------------------------------------------===// +// +// 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 _BTree { + /// Filters a B-Tree on a predicate, returning a new tree. + /// + /// - Complexity: O(`n log n`) where `n` is the number of key-value pairs in the + /// sorted dictionary. + @inlinable + @inline(__always) + public func filter( + _ isIncluded: (Element) throws -> Bool + ) rethrows -> _BTree { + var builder = Builder() + for element in self where try isIncluded(element) { + builder.append(element) + } + return builder.finish() + } + + + // MARK: Last Removal + + /// Removes the first element of a tree, if it exists. + /// + /// - Returns: The moved last 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 + } + + @inlinable + @inline(__always) + @discardableResult + public mutating func removeLast() -> Element { + if let value = self.popLast() { + return value + } else { + preconditionFailure("Can't remove last element from an empty collection") + } + } + + @inlinable + @inline(__always) + public mutating func removeLast(_ k: Int) { + assert(0 <= k && k < self.count, "Can't remove more items from a collection than it contains") + for _ in 0.. Element? { + invalidateIndices() + + if self.count == 0 { return nil } + + let removedElement = self.root.update { $0.popFirstElement() } + self._balanceRoot() + return removedElement + } + + @inlinable + @inline(__always) + @discardableResult + public mutating func removeFirst() -> Element { + if let value = self.popFirst() { + return value + } else { + preconditionFailure("Can't remove first element from an empty collection") + } + } + + @inlinable + @inline(__always) + public mutating func removeFirst(_ k: Int) { + assert(0 <= k && k < self.count, "Can't remove more items from a collection than it contains") + for _ in 0.. Element { + invalidateIndices() + guard index != endIndex else { preconditionFailure("Index out of bounds.") } + return self.remove(atOffset: index.offset) + } + + // MARK: Bulk Removal + @inlinable + @inline(__always) + internal mutating func removeAll() { + invalidateIndices() + // TODO: potentially use empty storage class. + self.root = _Node(withCapacity: _BTree.defaultLeafCapacity, isLeaf: true) + } + + /// Removes the elements in the specified subrange from the collection. + @inlinable + internal mutating func removeSubrange(_ bounds: Range) { + guard bounds.lowerBound != endIndex else { preconditionFailure("Index out of bounds.") } + guard bounds.upperBound != endIndex else { preconditionFailure("Index out of bounds.") } + + let rangeSize = self.distance(from: bounds.lowerBound, to: bounds.upperBound) + let startOffset = bounds.lowerBound.offset + + for _ in 0.. Void) rethrows { + func loop(node: Unmanaged) throws { + try node._withUnsafeGuaranteedRef { storage in + try storage.read { handle in + for i in 0..] + + /// 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 + + if _slowPath(self.tree.isEmpty || index.slot == -1) { + self.slots = [] + self.path = [] + return + } + + self.slots = [] + for d in 0.. Element? { + // Check slot sentinel value for end of tree. + if _slowPath(path.isEmpty) { + return nil + } + + let element = path[path.count - 1]._withUnsafeGuaranteedRef { + $0.read { $0[elementAt: Int(slots[slots.count - 1])] } + } + + path[path.count - 1]._withUnsafeGuaranteedRef { + $0.read({ handle in + self._advanceState(withLeaf: handle) + }) + } + + return element + } + } + + @inlinable + internal func makeIterator() -> Iterator { + return Iterator(forTree: self) + } +} diff --git a/Sources/SortedCollections/BTree/_BTree+SubSequence.swift b/Sources/SortedCollections/BTree/_BTree+SubSequence.swift new file mode 100644 index 000000000..52a2ce1e8 --- /dev/null +++ b/Sources/SortedCollections/BTree/_BTree+SubSequence.swift @@ -0,0 +1,187 @@ +//===----------------------------------------------------------------------===// +// +// 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 _BTree { + @usableFromInline + internal struct SubSequence { + @usableFromInline + internal let _base: _BTree + + @usableFromInline + internal var _startIndex: Index + + @usableFromInline + internal var _endIndex: Index + + @inlinable + @inline(__always) + internal init(base: _BTree, bounds: Range) { + self._base = base + self._startIndex = bounds.lowerBound + self._endIndex = bounds.upperBound + } + + /// The underlying collection of the subsequence. + @inlinable + @inline(__always) + internal var base: _BTree { _base } + } +} + +extension _BTree.SubSequence: Sequence { + @usableFromInline + internal typealias Element = _BTree.Element + + + @usableFromInline + internal struct Iterator: IteratorProtocol { + @usableFromInline + internal typealias Element = SubSequence.Element + + @usableFromInline + internal var _iterator: _BTree.Iterator + + @usableFromInline + internal var distanceRemaining: Int + + @inlinable + @inline(__always) + internal init(_iterator: _BTree.Iterator, distance: Int) { + self._iterator = _iterator + self.distanceRemaining = distance + } + + @inlinable + @inline(__always) + internal mutating func next() -> Element? { + if distanceRemaining == 0 { + return nil + } else { + distanceRemaining -= 1 + return _iterator.next() + } + } + } + + @inlinable + @inline(__always) + internal func makeIterator() -> Iterator { + let it = _BTree.Iterator(forTree: _base, startingAt: _startIndex) + let distance = _base.distance(from: _startIndex, to: _endIndex) + return Iterator(_iterator: it, distance: distance) + } +} + +extension _BTree.SubSequence: BidirectionalCollection { + @usableFromInline + internal typealias Index = _BTree.Index + + @usableFromInline + internal typealias SubSequence = Self + + + @inlinable + @inline(__always) + internal var startIndex: Index { _startIndex } + + @inlinable + @inline(__always) + internal var endIndex: Index { _endIndex } + + @inlinable + @inline(__always) + internal var count: Int { _base.distance(from: _startIndex, to: _endIndex) } + + @inlinable + @inline(__always) + internal func distance(from start: Index, to end: Index) -> Int { + _base.distance(from: start, to: end) + } + + @inlinable + @inline(__always) + internal func index(before i: Index) -> Index { + _base.index(before: i) + } + + @inlinable + @inline(__always) + internal func formIndex(before i: inout Index) { + _base.formIndex(before: &i) + } + + + @inlinable + @inline(__always) + internal func index(after i: Index) -> Index { + _base.index(after: i) + } + + @inlinable + @inline(__always) + internal func formIndex(after i: inout Index) { + _base.formIndex(after: &i) + } + + @inlinable + @inline(__always) + internal func index(_ i: Index, offsetBy distance: Int) -> Index { + _base.index(i, offsetBy: distance) + } + + @inlinable + @inline(__always) + internal func formIndex(_ i: inout Index, offsetBy distance: Int) { + _base.formIndex(&i, offsetBy: distance) + } + + @inlinable + @inline(__always) + internal func index(_ i: Index, offsetBy distance: Int, limitedBy limit: Index) -> Index? { + _base.index(i, offsetBy: distance, limitedBy: limit) + } + + @inlinable + @inline(__always) + internal func formIndex(_ i: inout Index, offsetBy distance: Int, limitedBy limit: Self.Index) -> Bool { + _base.formIndex(&i, offsetBy: distance, limitedBy: limit) + } + + + @inlinable + @inline(__always) + internal subscript(position: Index) -> Element { + _failEarlyRangeCheck(position, bounds: startIndex..) -> SubSequence { + _failEarlyRangeCheck(bounds, bounds: startIndex..) { + _base._failEarlyRangeCheck(index, bounds: bounds) + } + + @inlinable + @inline(__always) + public func _failEarlyRangeCheck(_ range: Range, bounds: Range) { + _base._failEarlyRangeCheck(range, bounds: bounds) + } +} + +// TODO: implement partial RangeReplaceableCollection methods diff --git a/Sources/SortedCollections/BTree/_BTree+UnsafeCursor.swift b/Sources/SortedCollections/BTree/_BTree+UnsafeCursor.swift new file mode 100644 index 000000000..14090c77c --- /dev/null +++ b/Sources/SortedCollections/BTree/_BTree+UnsafeCursor.swift @@ -0,0 +1,504 @@ +//===----------------------------------------------------------------------===// +// +// 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 _BTree { + /// A mutable cursor to an element of the B-Tree represented as a path. + /// + /// Cursors consume the original tree on which they were created. It is undefined behavior to operate or + /// read a tree on which a cursor was created. A cursor strongly references the tree on which it was + /// operating on. Once operations with a cursor are finished, the tree can be restored using + /// `_BTree.UnsafeCursor.apply(to:)`. + /// + /// This is a heavier alternative to ``_BTree.Index``, however this allows mutable operations to be + /// efficiently performed. + /// + /// - Warning: It is invalid to operate on a tree while a cursor to it exists. + /// - Warning: the tree root must remain alive for the entire lifetime of a cursor otherwise bad things + /// may occur. + @usableFromInline + internal struct UnsafeCursor { + @usableFromInline + internal typealias Path = _FixedSizeArray> + + /// This property is what takes ownership of the tree during the lifetime of the cursor. Once the cursor + /// is consumed, it is set to nil and it is invalid to use the cursor. + @usableFromInline + internal var _root: Node.Storage? + + + /// Position of each of the parent nodes in their parents, including the bottom-most node. + /// + /// In the following tree, a cursor to `7` would contain `[1, 0]`. Referring to the child at index 1 + /// and its element at index 0. + /// + /// ┌─┐ + /// │5│ + /// ┌─┴─┴─┐ + /// │ │ + /// ┌─┼─┐ ┌─┼─┐ + /// │1│3│ │7│9│ + /// └─┴─┘ └─┴─┘ + @usableFromInline + internal var slots: _FixedSizeArray<_BTree.Slot> + + /// This stores a list of the nodes from top-to-bottom. + @usableFromInline + internal var path: Path + + /// Bottom most node that the index point to. + + /// The depth at which the last instance of sequential unique nodes starting at the root was found. + /// + /// In the following path where 'U' denotes a unique node, and 'S' denotes a shared node. The value + /// of this parameter would be '1' indicating the second level of the tre. + /// + /// ┌─┐ + /// │U├─┐ + /// └─┘ ├─┐ + /// │U├─┐ + /// └─┘ ├─┐ + /// ▲ │S├─┐ + /// │ └─┘ ├─┐ + /// │ │U│ + /// └─┘ + /// + /// This is notable for CoW as all values below it would need to be duplicated. Updating this to be as + /// high as accurately possible ensures there are no unnecessary copies made. + @usableFromInline + internal var lastUniqueDepth: Int + + @inlinable + @inline(__always) + internal init( + root: Node.Storage, + slots: _FixedSizeArray<_BTree.Slot>, + path: Path, + lastUniqueDepth: Int + ) { + // Slots and path should be non-empty + assert(slots.depth >= 1, "Invalid tree cursor.") + assert(path.depth >= 1, "Invalid tree cursor.") + + self._root = root + self.slots = slots + self.path = path + self.lastUniqueDepth = lastUniqueDepth + } + + // MARK: Internal Checks + /// Check that this cursor is still valid + /// + /// Every member that operates on the element of the cursor must start by calling this function. + /// + /// Note that this is a noop in release builds. + @inlinable + @inline(__always) + internal func assertValid() { + #if COLLECTIONS_INTERNAL_CHECKS + assert(self._root != nil, + "Attempt to operate on an element using an invalid cursor.") + #endif + } + + // MARK: Core Cursor Operations + /// Finishes operating on a cursor and restores a tree + @inlinable + @inline(__always) + internal mutating func apply(to tree: inout _BTree) { + assertValid() + assert(tree.root._storage == nil, "Must apply to same tree as original.") + swap(&tree.root._storage, &self._root) + tree.checkInvariants() + } + + /// Declares that the cursor is completely unique + @inlinable + @inline(__always) + internal mutating func _declareUnique() { + self.lastUniqueDepth = Int(path.depth) + } + + /// Operators on a handle of the node + /// - Warning: Ensure this is never called on an endIndex. + @inlinable + @inline(__always) + internal mutating func readCurrentNode( + _ body: (Node.UnsafeHandle, Int) throws -> R + ) rethrows -> R { + assertValid() + + let slot = Int(slots[slots.depth - 1]) + return try path[path.depth - 1]._withUnsafeGuaranteedRef { + try $0.read({ try body($0, slot) }) + } + } + + /// Updates the node at a given depth. + /// - Warning: this does not perform CoW checks + @inlinable + @inline(__always) + internal mutating func updateNode( + atDepth depth: Int8, + _ body: (Node.UnsafeHandle, Int) throws -> R + ) rethrows -> (node: Node, result: R) { + assertValid() + + let slot = Int(slots[depth]) + let isOnUniquePath = depth <= lastUniqueDepth + + return try path[depth]._withUnsafeGuaranteedRef { storage in + if isOnUniquePath { + let result = try storage.updateGuaranteedUnique({ try body($0, slot) }) + return (Node(storage), result) + } else { + let storage = storage.copy() + path[depth] = .passUnretained(storage) + let result = try storage.updateGuaranteedUnique({ try body($0, slot) }) + return (Node(storage), result) + } + } + } + + + // MARK: Mutations with the Cursor + /// Operates on a handle of the node. + /// + /// This MUST be followed by an operation which consumes the cursor and returns the new tree, as + /// this only performs CoW on the bottom-most level + /// + /// - Parameter body: Updating callback to run on a unique instance of the node containing the + /// cursor's target. + /// - Returns: The body's return value + /// - Complexity: O(`log n`) if non-unique. O(`1`) if unique. + @inlinable + @inline(__always) + internal mutating func updateCurrentNode( + _ body: (Node.UnsafeHandle, Int) throws -> R + ) rethrows -> R { + assertValid() + defer { self._declareUnique() } + + // Update the bottom-most node + var (node, result) = try self.updateNode(atDepth: path.depth - 1, body) + + // Start the node above the bottom-most node, and propogate up the change + var depth = path.depth - 2 + while depth >= 0 { + if depth > lastUniqueDepth { + // If we're on a + let (newNode, _) = self.updateNode(atDepth: depth) { (handle, slot) in + _ = handle.exchangeChild(atSlot: slot, with: node) + } + + node = newNode + } else if depth == lastUniqueDepth { + // The node directly above the first shared node, we can update its + // child without performing uniqueness checks since it's guaranteed + // to be unique. + let child = node + let slot = Int(slots[depth]) + + path[depth]._withUnsafeGuaranteedRef { storage in + storage.updateGuaranteedUnique { handle in + _ = handle.exchangeChild(atSlot: slot, with: child) + } + } + + return result + } else { + // depth < lastUniqueDepth + + // In this case we don't need to traverse to the root since we know + // it remains the same. + return result + } + + depth -= 1 + } + + self._root = node.storage + return result + } + + /// Moves the value from the cursor's position + @inlinable + @inline(__always) + internal mutating func moveValue() -> Value { + guard Node.hasValues else { return Node.dummyValue } + + return self.updateCurrentNode { handle, slot in + handle.pointerToValue(atSlot: slot).move() + } + } + + /// Initializes a value for a cursor that points to an element that has a hole for its value. + @inlinable + @inline(__always) + internal mutating func initializeValue(to value: Value) { + guard Node.hasValues else { return } + + self.updateCurrentNode { handle, slot in + handle.pointerToValue(atSlot: slot).initialize(to: value) + } + } + + /// Inserts a key-value pair at a position within the tree. + /// + /// Invalidates the cursor. Returns a splinter object which owning the new tree. + /// + /// - Parameters: + /// - element: A new key-value element to insert. + /// - capacity: Capacity of new internal nodes created during insertion. + /// - Returns: The new root object which may equal in identity to the previous one. + /// - Warning: Doesn not check sortedness invariant + /// - Complexity: O(`log n`). Ascends the tree once + @inlinable + internal mutating func insertElement( + _ element: Node.Element, + capacity: Int + ) { + assertValid() + defer { self._declareUnique() } + + var (node, splinter) = self.updateNode(atDepth: path.depth - 1) { handle, slot in + handle.insertElement(element, withRightChild: nil, atSlot: slot) + } + + // Start the node above the bottom-most node, and propogate up the change + var depth = path.depth - 2 + while depth >= 0 { + let (newNode, _) = self.updateNode(atDepth: depth) { (handle, slot) in + handle.exchangeChild(atSlot: slot, with: node) + + if let lastSplinter = splinter { + splinter = handle.insertSplinter(lastSplinter, atSlot: slot) + } else { + handle.subtreeCount += 1 + } + } + + node = newNode + depth -= 1 + } + + if let splinter = splinter { + let newRoot = splinter.toNode(leftChild: node, capacity: capacity) + self._root = newRoot.storage + } else { + self._root = node.storage + } + } + + /// Removes the element at a cursor. + /// + /// - Parameter hasValueHole: Whether the value has been moved out of the node. + /// - Complexity: O(`log n`). Ascends the tree once. + @inlinable + internal mutating func removeElement(hasValueHole: Bool = false) { + assertValid() + defer { self._declareUnique() } + + var (node, _) = self.updateNode( + atDepth: path.depth - 1 + ) { handle, slot in + if handle.isLeaf { + // Deletion within a leaf + // removeElement(atSlot:) automatically adjusts node counts. + if hasValueHole { + handle.removeElementWithoutValue(atSlot: slot) + } else { + handle.removeElement(atSlot: slot) + } + } else { + // Deletion within an internal node + + // Swap with the predecessor + let predecessor = + handle[childAt: slot].update { $0.popLastElement() } + + // Reduce the element count. + handle.subtreeCount -= 1 + + // Replace the current element with the predecessor. + if hasValueHole { + _ = handle.pointerToKey(atSlot: slot).move() + handle.initializeElement(atSlot: slot, to: predecessor) + } else { + handle.exchangeElement(atSlot: slot, with: predecessor) + } + + // Balance the predecessor child slot, as the pop operation may have + // brought it out of balance. + handle.balance(atSlot: slot) + } + } + + // Balance the parents + var depth = path.depth - 2 + while depth >= 0 { + var (newNode, _) = self.updateNode(atDepth: depth) { (handle, slot) in + handle.exchangeChild(atSlot: slot, with: node) + handle.subtreeCount -= 1 + handle.balance(atSlot: slot) + } + + if depth == 0 && newNode.read({ $0.elementCount == 0 && !$0.isLeaf }) { + // If the root has no elements, we won't + newNode.update(isUnique: true) { $0.drop() } + } else { + node = newNode + } + depth -= 1 + } + + self._root = node.storage + } + } + + /// Obtains a cursor to a given element in the tree. + /// + /// This 'consumes' the tree, however it expects the callee to retain the root of the tree for the duration of + /// the cursors lifetime. + /// + /// - Parameter key: The key to search for + /// - Returns: A cursor to the key or where the key should be inserted. + /// - Complexity: O(`log n`) + @inlinable + internal mutating func takeCursor(at index: Index) -> UnsafeCursor { + var slots = index.childSlots + slots.append(UInt16(index.slot)) + + // Initialize parents with some dummy value filling it. + var parents = + UnsafeCursor.Path(repeating: .passUnretained(self.root.storage)) + + var ownedRoot: Node.Storage + do { + var tempRoot: Node.Storage? = nil + swap(&tempRoot, &self.root._storage) + ownedRoot = tempRoot.unsafelyUnwrapped + } + + var node: Unmanaged = .passUnretained(ownedRoot) + + // The depth containing the first instance of a shared + var lastUniqueDepth = isKnownUniquelyReferenced(&ownedRoot) ? 0 : -1 + var isOnUniquePath = isKnownUniquelyReferenced(&ownedRoot) + + for d in 0.. (cursor: UnsafeCursor, found: Bool) { + var slots = Index.Offsets(repeating: 0) + + // Initialize parents with some dummy value filling it. + var parents = + UnsafeCursor.Path(repeating: .passUnretained(self.root.storage)) + + var ownedRoot: Node.Storage + do { + var tempRoot: Node.Storage? = nil + swap(&tempRoot, &self.root._storage) + ownedRoot = tempRoot.unsafelyUnwrapped + } + + var node: Unmanaged = .passUnretained(ownedRoot) + + // Initialize slot to some dummy value. + var slot = -1 + var found: Bool = false + + // The depth containing the first instance of a shared + var lastUniqueDepth = isKnownUniquelyReferenced(&ownedRoot) ? 0 : -1 + var isOnUniquePath = isKnownUniquelyReferenced(&ownedRoot) + + while true { + let shouldStop: Bool = node._withUnsafeGuaranteedRef { storage in + storage.read { handle in + slot = handle.startSlot(forKey: key) + + if slot < handle.elementCount && handle[keyAt: slot] == key { + found = true + return true + } else { + if handle.isLeaf { + return true + } else { + parents.append(node) + slots.append(UInt16(slot)) + + node = .passUnretained(handle[childAt: slot].storage) + if isOnUniquePath && handle.isChildUnique(atSlot: slot) { + lastUniqueDepth += 1 + } else { + isOnUniquePath = false + } + return false + } + } + } + } + + if shouldStop { break } + } + + assert(slot != -1, "B-Tree sanity check fail.") + + parents.append(node) + slots.append(UInt16(slot)) + + let cursor = UnsafeCursor( + root: ownedRoot, + slots: slots, + path: parents, + lastUniqueDepth: lastUniqueDepth + ) + + return (cursor, found) + } +} diff --git a/Sources/SortedCollections/BTree/_BTree.Builder.swift b/Sources/SortedCollections/BTree/_BTree.Builder.swift new file mode 100644 index 000000000..7104e8d85 --- /dev/null +++ b/Sources/SortedCollections/BTree/_BTree.Builder.swift @@ -0,0 +1,370 @@ +//===----------------------------------------------------------------------===// +// +// 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 _BTree { + /// Provides an interface for efficiently constructing a filled B-Tree from sorted data. + /// + /// A builder supports duplicate keys, in which case they are inserted in the same order they are recieved. + /// However, is the `deduplicating` parameter is passed as `true`, operations will silently drop + /// duplicates. + /// + /// This type has a few advantages when constructing a B-Tree over other approaches such as manually + /// inserting each element or using a cursor: + /// + /// This works by maintaing a list of saplings and a view of the node currently being modified. For example + /// the following tree: + /// + /// ┌─┐ + /// │D│ + /// ┌───┴─┴───┐ + /// │ │ + /// ┌┴┐ ┌┴┐ + /// │B│ │F│ + /// ┌─┴─┴─┐ ┌─┴─┴─┐ + /// │ │ │ │ + /// ┌┴┐ ┌┴┐ ┌┴┐ ┌┴┐ + /// │A│ │C│ │E│ │G│ + /// └─┘ └─┘ └─┘ └─┘ + /// + /// Would be represented in the following state: + /// + /// ┌─┐ + /// Seedling: │G│ + /// └─┘ + /// + /// ┌─┐ + /// │B│ ┌─┐ + /// Saplings: ┌─┴─┴─┐ │E│ + /// │ │ └─┘ + /// ┌┴┐ ┌┴┐ + /// │A│ │C│ + /// └─┘ └─┘ + /// + /// ┌─┐ ┌─┐ + /// Seperators: │D│ │F│ + /// └─┘ └─┘ + /// + /// While the diagrams above represent a binary-tree, the representation of a B-Tree in the builder is + /// directly analogous to this. By representing the state this way. Append operations can be efficiently + /// performed, and the tree can also be efficiently reconstructed. + /// + /// Appending works by filling in a seedling, once a seedling is full, and an associated seperator has been + /// provided, the seedling-seperator pair can be appended to the stack. + @usableFromInline + internal struct Builder { + @usableFromInline + enum State { + /// The builder needs to add a seperator to the node + case addingSeperator + + /// The builder needs to try to append to the seedling node. + case appendingToSeedling + } + + @usableFromInline + internal var _saplings: [Node] + + @usableFromInline + internal var _seperators: [Element] + + @usableFromInline + internal var _seedling: Node? + + @inlinable + @inline(__always) + internal var seedling: Node { + get { + assert(_seedling != nil, + "Simultaneous access or access on consumed builder.") + return _seedling.unsafelyUnwrapped + } + _modify { + assert(_seedling != nil, + "Simultaneous mutable access or mutable access on consumed builder.") + var value = _seedling.unsafelyUnwrapped + _seedling = nil + defer { _seedling = value } + yield &value + } + } + + @usableFromInline + internal var state: State + + @usableFromInline + internal let leafCapacity: Int + + @usableFromInline + internal let internalCapacity: Int + + @usableFromInline + internal let deduplicating: Bool + + @usableFromInline + internal var lastKey: Key? + + /// Creates a new B-Tree builder with default capacities + /// - Parameter deduplicating: Whether duplicates should be removed. + @inlinable + @inline(__always) + internal init(deduplicating: Bool = false) { + self.init( + deduplicating: deduplicating, + leafCapacity: _BTree.defaultLeafCapacity, + internalCapacity: _BTree.defaultInternalCapacity + ) + } + + /// Creates a new B-Tree builder with a custom uniform capacity configuration + /// - Parameters: + /// - deduplicating: Whether duplicates should be removed. + /// - capacity: The amount of elements per node. + @inlinable + @inline(__always) + internal init(deduplicating: Bool = false, capacity: Int) { + self.init( + deduplicating: deduplicating, + leafCapacity: capacity, + internalCapacity: capacity + ) + } + + /// Creates a new B-Tree builder with a custom capacity configuration + /// - Parameters: + /// - deduplicating: Whether duplicates should be removed. + /// - leafCapacity: The amount of elements per leaf node. + /// - internalCapacity: The amount of elements per internal node. + @inlinable + @inline(__always) + internal init( + deduplicating: Bool = false, + leafCapacity: Int, + internalCapacity: Int + ) { + assert(leafCapacity > 1 && internalCapacity > 1, + "Capacity must be greater than one") + + self._saplings = [] + self._seperators = [] + self.state = .appendingToSeedling + self._seedling = Node(withCapacity: leafCapacity, isLeaf: true) + self.leafCapacity = leafCapacity + self.internalCapacity = internalCapacity + self.deduplicating = deduplicating + self.lastKey = nil + } + + /// Pops a sapling and it's associated seperator + @inlinable + @inline(__always) + internal mutating func popSapling() + -> (leftNode: Node, seperator: Element)? { + return _saplings.isEmpty ? nil : ( + leftNode: _saplings.removeLast(), + seperator: _seperators.removeLast() + ) + } + + /// Appends a sapling with an associated seperator + @inlinable + @inline(__always) + internal mutating func appendSapling( + _ sapling: __owned Node, + seperatedBy seperator: Element + ) { + _saplings.append(sapling) + _seperators.append(seperator) + } + + /// Appends a sequence of sorted values to the tree + @inlinable + @inline(__always) + internal mutating func append( + contentsOf sequence: S + ) where S.Element == Element { + for element in sequence { + self.append(element) + } + } + + /// Appends a new element to the tree + /// - Parameter element: Element which is after all previous elements in sorted order. + @inlinable + internal mutating func append(_ element: __owned Element) { + assert(lastKey == nil || lastKey! <= element.key, + "New element must be non-decreasing.") + defer { lastKey = element.key } + if deduplicating { + if let lastKey = lastKey { + if lastKey == element.key { return } + } + } + + switch state { + case .addingSeperator: + completeSeedling(withSeperator: element) + state = .appendingToSeedling + + case .appendingToSeedling: + let isFull: Bool = seedling.update { handle in + handle.appendElement(element) + return handle.isFull + } + + if _slowPath(isFull) { + state = .addingSeperator + } + } + } + + + + /// Declares that the current seedling is finished with insertion and creates a new seedling to + /// further operate on. + @inlinable + internal mutating func completeSeedling( + withSeperator newSeperator: __owned Element + ) { + var sapling = Node(withCapacity: leafCapacity, isLeaf: true) + swap(&sapling, &self.seedling) + + // Prepare a new sapling to insert. + // There are a few invariants we're thinking about here: + // - Leaf nodes are coming in fully filled. We can treat them as atomic + // bits + // - The stack has saplings of decreasing depth. + // - Saplings on the stack are completely filled except for their roots. + if case (var previousSapling, let seperator)? = self.popSapling() { + let saplingDepth = sapling.storage.header.depth + let previousSaplingDepth = previousSapling.storage.header.depth + let previousSaplingIsFull = previousSapling.read({ $0.isFull }) + + assert(previousSaplingDepth >= saplingDepth, + "Builder invariant failure.") + + if saplingDepth == previousSaplingDepth && previousSaplingIsFull { + // This is when two nodes are full: + // + // ┌───┐ ┌───┐ + // │ A │ │ C │ + // └───┘ └───┘ + // ▲ ▲ + // │ │ + // previousSapling sapling + // + // We then use the seperator (B) to transform this into a subtree of a + // depth increase: + // ┌───┐ + // │ B │ ◄─── sapling + // ┌┴───┴┐ + // │ │ + // ┌─┴─┐ ┌─┴─┐ + // │ A │ │ C │ + // └───┘ └───┘ + // If the sapling is full. We create a splinter. This is when the + // depth of our B-Tree increases + sapling = _Node( + leftChild: previousSapling, + seperator: seperator, + rightChild: sapling, + capacity: internalCapacity + ) + } else if saplingDepth + 1 == previousSaplingDepth && !previousSaplingIsFull { + // This is when we can append the node with the seperator: + // + // ┌───┐ + // │ B │ ◄─ previousSapling + // ┌┴───┴┐ + // │ │ + // ┌─┴─┐ ┌─┴─┐ ┌───┐ + // │ A │ │ C │ │ E │ ◄─ sapling + // └───┘ └───┘ └───┘ + // + // We then use the seperator (D) to append this to previousSapling. + // ┌────┬───┐ + // │ B │ D │ ◄─ sapling + // ┌┴────┼───┴┐ + // │ │ │ + // ┌─┴─┐ ┌─┴─┐ ┌┴──┐ + // │ A │ │ C │ │ E │ + // └───┘ └───┘ └───┘ + previousSapling.update { + $0.appendElement(seperator, withRightChild: sapling) + } + sapling = previousSapling + } else { + // In this case, we need to work on creating a new sapling. Say we + // have: + // + // ┌────┬───┐ + // │ B │ D │ ◄─ previousSapling + // ┌┴────┼───┴┐ + // │ │ │ + // ┌─┴─┐ ┌─┴─┐ ┌┴──┐ ┌───┐ + // │ A │ │ C │ │ E │ │ G │ ◄─ sapling + // └───┘ └───┘ └───┘ └───┘ + // + // Where previousSapling is full. We'll commit sapling and keep + // working on it until it is of the same depth as `previousSapling`. + // Once it is the same depth, we can join the nodes. + // + // The goal is once we have a full tree of equal depth: + // + // ┌────┬───┐ ┌────┬───┐ + // │ B │ D │ │ H │ J │ + // ┌┴────┼───┴┐ ┌┴────┼───┴┐ + // │ │ │ │ │ │ + // ┌─┴─┐ ┌─┴─┐ ┌┴──┐ ┌─┴─┐ ┌─┴─┐ ┌┴──┐ + // │ A │ │ C │ │ E │ │ G │ │ I │ │ K │ + // └───┘ └───┘ └───┘ └───┘ └───┘ └───┘ + // + // We can string them together using the previous cases. + self.appendSapling(previousSapling, seperatedBy: seperator) + } + } + + self.appendSapling(sapling, seperatedBy: newSeperator) + } + + /// Finishes building a tree. + /// + /// This consumes the builder and it is no longer valid to operate on after this. + /// + /// - Returns: A usable, fully-filled B-Tree + @inlinable + internal mutating func finish() -> _BTree { + var root: Node = seedling + _seedling = nil + + while case (var sapling, let seperator)? = self.popSapling() { + root = _Node.join( + &sapling, + with: &root, + seperatedBy: seperator, + capacity: internalCapacity + ) + } + + let tree = _BTree(rootedAt: root, internalCapacity: internalCapacity) + tree.checkInvariants() + return tree + } + } +} + +extension _BTree.Builder where Value == Void { + /// Appends a value to a B-Tree builder without values. + @inlinable + @inline(__always) + internal mutating func append(_ key: __owned Key) { + self.append((key, ())) + } +} diff --git a/Sources/SortedCollections/BTree/_BTree.Index.swift b/Sources/SortedCollections/BTree/_BTree.Index.swift new file mode 100644 index 000000000..31b43fc45 --- /dev/null +++ b/Sources/SortedCollections/BTree/_BTree.Index.swift @@ -0,0 +1,142 @@ +//===----------------------------------------------------------------------===// +// +// 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 _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 { + /// A fixed-size array large enough to represent all offsets within a B-Tree index + @usableFromInline + internal typealias Offsets = _FixedSizeArray + + /// The position of each of the parent nodes in their parents. The path's depth + /// is offsets.count + 1 + @usableFromInline + internal var childSlots: Offsets + + /// The bottom most node that the index point to + /// + /// This is equal to `root` to indicate the `endIndex` + @usableFromInline + internal var node: Unmanaged + + /// The slot within the bottom most node which the index points to. + /// + /// This is equal to `-1` to indicate the `endIndex` + @usableFromInline + internal var slot: Int + + /// The absolute offset of the path's element in the entire tree. + @usableFromInline + internal var offset: Int + + /// The tree that this index references. + @usableFromInline + internal weak var root: Node.Storage? + + /// The age of the tree when this index was captures. + @usableFromInline + internal let version: Int + + /// Creates an index represented as a sequence of nodes to an element. + /// - Parameters: + /// - node: The node to which this index points. + /// - slot: The specific slot within node where the path points + /// - childSlots: The children's offsets for this path. + /// - index: The absolute offset of this path's element in the tree. + /// - tree: The tree of this index. + @inlinable + @inline(__always) + internal init( + node: Unmanaged, + slot: Int, + childSlots: Offsets, + offset: Int, + forTree tree: _BTree + ) { + self.node = node + self.slot = slot + self.childSlots = childSlots + self.offset = offset + + self.root = tree.root.storage + self.version = tree.version + } + + + // MARK: Index Validation + /// Ensures the precondition that the index is valid for a given dictionary. + @inlinable + @inline(__always) + internal func ensureValid(forTree tree: _BTree) { + precondition( + self.root === tree.root.storage && self.version == tree.version, + "Attempt to use an invalid index.") + } + + /// Ensures the precondition that the index is valid for use with another index + @inlinable + @inline(__always) + internal func ensureValid(with index: Index) { + precondition( + self.root != nil && + self.root === index.root && + self.version == index.version, + "Attempt to use an invalid indices.") + } + } +} + +// MARK: Index Read Operations +extension _BTree.Index { + /// Operators on a handle of the node + /// - Warning: Ensure this is never called on an endIndex. + @inlinable + @inline(__always) + internal func readNode( + _ body: (_BTree.Node.UnsafeHandle) throws -> R + ) rethrows -> R { + assert(self.slot != -1, "Invalid operation to read end index.") + return try self.node._withUnsafeGuaranteedRef { try $0.read(body) } + } + + /// Gets the element the path points to. + @inlinable + @inline(__always) + internal var element: _BTree.Element { + assert(self.slot != -1, "Cannot dereference out-of-bounds slot.") + return self.readNode { $0[elementAt: self.slot] } + } +} + +// MARK: Comparable +extension _BTree.Index: Comparable { + /// Returns true if two indices are identical (point to the same node). + /// - Precondition: expects both indices are from the same B-Tree. + /// - Complexity: O(1) + @inlinable + @inline(__always) + internal static func ==(lhs: Self, rhs: Self) -> Bool { + return lhs.offset == rhs.offset + } + + /// Returns true if the first path points to an element before the second path + /// - Precondition: expects both indices are from the same B-Tree. + /// - Complexity: O(1) + @inlinable + @inline(__always) + internal static func <(lhs: Self, rhs: Self) -> Bool { + return lhs.offset < rhs.offset + } +} diff --git a/Sources/SortedCollections/BTree/_BTree.swift b/Sources/SortedCollections/BTree/_BTree.swift new file mode 100644 index 000000000..488265676 --- /dev/null +++ b/Sources/SortedCollections/BTree/_BTree.swift @@ -0,0 +1,578 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +/// A bidirectional collection representing a BTree which efficiently stores its +/// elements in sorted order and maintains roughly `O(log count)` +/// performance for most operations. +/// +/// - Warning: Indexing operations on a BTree are unchecked. ``_BTree.Index`` +/// offers `ensureValid(for:)` methods to validate indices for use in higher-level +/// collections. +@usableFromInline +internal struct _BTree { + + /// Recommended node size of a given B-Tree + @inlinable + @inline(__always) + internal static var defaultInternalCapacity: Int { + #if DEBUG + return 4 + #else + let capacityInBytes = 128 + return Swift.min(16, capacityInBytes / MemoryLayout.stride) + #endif + } + + /// Recommended node size of a given B-Tree + @inlinable + @inline(__always) + internal static var defaultLeafCapacity: Int { + #if DEBUG + return 5 + #else + let capacityInBytes = 2000 + return Swift.min(16, capacityInBytes / MemoryLayout.stride) + #endif + } + + /// The element type of the collection. + @usableFromInline + internal typealias Element = (key: Key, value: Value) + + /// The type of each node in the tree + @usableFromInline + internal typealias Node = _Node + + /// A size large enough to represent any slot within a node + @usableFromInline + internal typealias Slot = UInt16 + + /// The underlying node behind this local BTree + @usableFromInline + internal var root: Node + + /// The capacity of each of the internal nodes + @usableFromInline + internal var internalCapacity: Int + + /// A metric to uniquely identify a given B-Tree's state. It is not + /// impossible for two B-Trees to have the same age by pure + /// coincidence. + @usableFromInline + internal var version: Int + + /// Creates a dummy B-Tree with no underlying node storage class. + /// + /// It is invalid and a serious error to ever attempt to read or write to such a B-Tree. + @inlinable + @inline(__always) + internal static var dummy: _BTree { + _BTree( + _rootedAtNode: _Node.dummy, + internalCapacity: 0, + version: 0 + ) + } + + /// Creates an empty B-Tree with an automatically determined optimal capacity. + @inlinable + @inline(__always) + internal init() { + let root = Node(withCapacity: _BTree.defaultLeafCapacity, isLeaf: true) + self.init(rootedAt: root, internalCapacity: _BTree.defaultInternalCapacity) + } + + /// Creates an empty B-Tree rooted at a specific node with a specified uniform capacity + /// - Parameter capacity: The key capacity of all nodes. + @inlinable + @inline(__always) + internal init(capacity: Int) { + self.init( + leafCapacity: capacity, + internalCapacity: capacity + ) + } + + /// Creates an empty B-Tree which creates node with specified capacities + /// - Parameters: + /// - leafCapacity: The capacity of the leaf nodes. This is the initial buffer used to allocate. + /// - internalCapacity: The capacity of the internal nodes. Generally prefered to be less than + /// `leafCapacity`. + @inlinable + @inline(__always) + internal init(leafCapacity: Int, internalCapacity: Int) { + self.init( + rootedAt: Node(withCapacity: leafCapacity, isLeaf: true), + internalCapacity: internalCapacity + ) + } + + /// Creates a B-Tree rooted at a specific nodey + /// - Parameters: + /// - root: The root node. + /// - internalCapacity: The key capacity of new internal nodes. + @inlinable + @inline(__always) + internal init(rootedAt root: Node, internalCapacity: Int) { + self.root = root + self.internalCapacity = internalCapacity + self.version = ObjectIdentifier(root.storage).hashValue + } + + @inlinable + @inline(__always) + internal init( + _rootedAtNode root: Node, + internalCapacity: Int, + version: Int + ) { + self.root = root + self.internalCapacity = internalCapacity + self.version = version + } +} + +// MARK: Mutating Operations +extension _BTree { + /// Invalidates the issued indices of the dictionary. Ensure this is + /// called for operations which mutate the SortedDictionary. + @inlinable + @inline(__always) + internal mutating func invalidateIndices() { + self.version &+= 1 + } + + /// Inserts an element into the BTree, or updates it if it already exists within the tree. If there are + /// multiple instances of the key in the tree, this may update any one. + /// + /// This operation is marginally more efficient on trees without duplicates, and may have + /// inconsistent results on trees with duplicates. + /// + /// - Parameters: + /// - value: The value to set corresponding to the key. + /// - key: The key to search for. + /// - updatingKey: If the key is found, whether it should be updated. + /// - Returns: If updated, the previous element. + /// - Complexity: O(`log n`) + @inlinable + @discardableResult + internal mutating func updateAnyValue( + _ value: Value, + forKey key: Key, + updatingKey: Bool = false + ) -> Element? { + invalidateIndices() + defer { self.checkInvariants() } + + let result = self.root.update { $0.updateAnyValue(value, forKey: key, updatingKey: updatingKey) } + switch result { + case let .updated(previousValue): + return previousValue + case let .splintered(splinter): + self.root = splinter.toNode( + leftChild: self.root, + capacity: self.internalCapacity + ) + default: break + } + + 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.elementCount == 0 && !$0.isLeaf }) { + let newRoot: Node = self.root.update { handle in + let newRoot = handle.moveChild(atSlot: 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 + /// that do not have duplicates. + /// + /// If the key is not found, the tree is not modified, although the version of the tree may change. + /// + /// - Parameter key: The key to remove in the tree + /// - Returns: The key-value pair which was removed. `nil` if not removed. + @inlinable + @discardableResult + internal mutating func removeAnyElement(forKey key: Key) -> Element? { + invalidateIndices() + defer { self.checkInvariants() } + + // TODO: It is possible that there is a single transient CoW copied made + // if the removal results in the CoW copied node being completely removed + // from the tree. + let removedElement = self.root.update { $0.removeAnyElement(forKey: key) } + + // Check if the tree height needs to be reduced + self._balanceRoot() + self.checkInvariants() + + return removedElement + } +} + +// MARK: Removal Opertations +extension _BTree { + /// 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 + @inline(__always) + @discardableResult + internal mutating func remove(atOffset offset: Int) -> Element { + invalidateIndices() + defer { self.checkInvariants() } + + let removedElement = self.root.update { $0.remove(at: offset) } + self._balanceRoot() + + return removedElement + } +} + +// MARK: Read Operations +extension _BTree { + /// Determines if a key exists within a tree. + /// + /// - Parameter key: The key to search for. + /// - Returns: Whether or not the key was found. + @inlinable + internal func contains(key: Key) -> Bool { + // the retain/release calls + // Retain + var node: Node? = self.root + + while let currentNode = node { + let found: Bool = currentNode.read { handle in + let slot = handle.startSlot(forKey: key) + + if slot < handle.elementCount && handle[keyAt: slot] == key { + return true + } else { + if handle.isLeaf { + node = nil + } else { + // Release + // Retain + node = handle[childAt: slot] + } + } + + return false + } + + if found { return true } + } + + 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 + /// for trees that do not have duplicates. + /// + /// - Parameter key: The key to search for + /// - Returns: `nil` if the key was not found. Otherwise, the previous value. + /// - Complexity: O(`log n`) + @inlinable + internal func findAnyValue(forKey key: Key) -> Value? { + var node: Unmanaged? = .passUnretained(self.root.storage) + + while let currentNode = node { + let value: Value? = currentNode._withUnsafeGuaranteedRef { + $0.read { handle in + let slot = handle.startSlot(forKey: key) + + if slot < handle.elementCount && handle[keyAt: slot] == key { + return handle[valueAt: slot] + } else { + if handle.isLeaf { + node = nil + } else { + node = .passUnretained(handle[childAt: slot].storage) + } + } + + return nil + } + } + + if let value = value { return value } + } + + return nil + } + + /// 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`. + @inlinable + internal func findAnyIndex(forKey key: Key) -> Index? { + var childSlots = Index.Offsets(repeating: 0) + var node: Unmanaged? = .passUnretained(self.root.storage) + var offset: Int = 0 + + while let currentNode = node { + let index: Index? = currentNode._withUnsafeGuaranteedRef { storage in + storage.read { handle in + let keySlot = handle.startSlot(forKey: key) + offset += keySlot + + if keySlot < handle.elementCount && handle[keyAt: keySlot] == key { + return Index( + node: .passUnretained(storage), + slot: keySlot, + childSlots: childSlots, + offset: offset, + forTree: self + ) + } else { + if handle.isLeaf { + node = nil + } else { + for i in 0.. Index { + assert(offset <= self.count, "Index out of bounds.") + + // Return nil path if at the end of the tree + if offset == self.count { + return self.endIndex + } + + var childSlots = Index.Offsets(repeating: 0) + + var node: Unmanaged = .passUnretained(self.root.storage) + var startIndex = 0 + + while true { + let index: Index? = node._withUnsafeGuaranteedRef { storage in + storage.read({ handle in + if handle.isLeaf { + return Index( + node: node, + slot: offset - startIndex, + childSlots: childSlots, + offset: offset, + forTree: self + ) + } + + for childSlot in 0.. Index { + var childSlots = Index.Offsets(repeating: 0) + var targetSlot: Int = 0 + var offset = 0 + + func search(in node: Node) -> Unmanaged? { + node.read({ handle in + let slot = handle.startSlot(forKey: key) + if slot < handle.elementCount { + if handle.isLeaf { + offset += slot + targetSlot = slot + return .passUnretained(node.storage) + } else { + // 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 targetChild = search(in: self.root) { + return Index( + node: targetChild, + slot: targetSlot, + childSlots: childSlots, + offset: offset, + forTree: self + ) + } else { + return endIndex + } + } + + /// Obtains the last index at which a value less than or equal to the key appears. + @inlinable + internal func lastIndex(forKey key: Key) -> Index { + var childSlots = Index.Offsets(repeating: 0) + var targetSlot: Int = 0 + var offset = 0 + + func search(in node: Node) -> Unmanaged? { + 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 handle.isLeaf { + offset += slot + targetSlot = slot + return .passUnretained(node.storage) + } else { + 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 + } + }) + } + + if let targetChild = search(in: self.root) { + return Index( + node: targetChild, + slot: targetSlot, + childSlots: childSlots, + offset: offset, + forTree: self + ) + } else { + return endIndex + } + } + +} + +// MARK: Immutable Operatoins +extension _BTree { + @inlinable + @inline(__always) + public func mapValues( + _ transform: (Value) throws -> T + ) rethrows -> _BTree { + let root = try _Node( + mappingFrom: self.root, + transform + ) + + return _BTree(rootedAt: root, internalCapacity: internalCapacity) + } +} diff --git a/Sources/SortedCollections/BTree/_FixedSizeArray.swift b/Sources/SortedCollections/BTree/_FixedSizeArray.swift new file mode 100644 index 000000000..51a6c652b --- /dev/null +++ b/Sources/SortedCollections/BTree/_FixedSizeArray.swift @@ -0,0 +1,133 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + + /// A stack-allocated deque of some values. +/// +/// The supports efficient removals from and insertions to the beginning and end. +/// +/// - Warning: This may hold strong references to objects after they +/// do not put non-trivial types in this. +@usableFromInline +internal struct _FixedSizeArray { + @inlinable + @inline(__always) + internal static var maxSize: Int8 { 16 } + + @usableFromInline + internal var depth: Int8 + + @usableFromInline + internal var values: ( + Element, Element, Element, Element, + Element, Element, Element, Element, + Element, Element, Element, Element, + Element, Element, Element, Element + ) + + @inlinable + @inline(__always) + internal init(repeating initialValue: Element, depth: Int8 = 0) { + self.depth = depth + self.values = ( + initialValue, initialValue, initialValue, initialValue, + initialValue, initialValue, initialValue, initialValue, + initialValue, initialValue, initialValue, initialValue, + initialValue, initialValue, initialValue, initialValue + ) + } + + /// Appends a value to the offset list + @inlinable + @inline(__always) + internal mutating func append(_ value: __owned Element) { + assert(depth < _FixedSizeArray.maxSize, + "Out of bounds access in fixed sized array.") + defer { self.depth &+= 1 } + self[self.depth] = value + } + + /// Pops a value from the end of offset list + @inlinable + @inline(__always) + internal mutating func pop() -> Element { + assert(depth > 0, "Cannot pop empty fixed sized array") + self.depth &-= 1 + return self[self.depth] + } + + /// If the fixed size array is empty + @inlinable + @inline(__always) + internal var isEmpty: Bool { depth == 0 } + + /// Refers to the last value in the list + @inlinable + @inline(__always) + internal var last: Element { + get { + assert(depth > 0, "Out of bounds access in fixed sized array") + return self[depth &- 1] + } + + _modify { + assert(depth > 0, "Out of bounds access in fixed sized array") + yield &self[depth &- 1] + } + } + + @inlinable + @inline(__always) + internal subscript(_ position: Int8) -> Element { + get { + assert(position <= depth && depth <= _FixedSizeArray.maxSize, + "Out of bounds access in fixed sized array.") + + return withUnsafeBytes(of: self.values) { values in + let p = values.baseAddress!.assumingMemoryBound(to: Element.self) + return p.advanced(by: Int(position)).pointee + } + } + + _modify { + assert(position <= depth && depth <= _FixedSizeArray.maxSize, + "Out of bounds access in fixed sized array.") + + let ptr: UnsafeMutablePointer = + withUnsafeMutableBytes(of: &self.values) { values in + let p = values.baseAddress!.assumingMemoryBound(to: Element.self) + return p.advanced(by: Int(position)) + } + + var value = ptr.move() + defer { ptr.initialize(to: value) } + yield &value + } + } +} + +#if DEBUG +extension _FixedSizeArray: CustomDebugStringConvertible { + @inlinable + internal var debugDescription: String { + var result = "[" + + for i in 0..( + _keyValuePairs keyValuePairs: C, + children: [_Node]? = nil, + capacity: Int + ) where C.Element == Element { + precondition(keyValuePairs.count <= capacity, "Too many key-value pairs.") + self.init(withCapacity: capacity, isLeaf: children == nil) + + self.update { handle in + let sortedKeyValuePairs = keyValuePairs.sorted(by: { $0.key < $1.key }) + + for (index, pair) in sortedKeyValuePairs.enumerated() { + let (key, value) = pair + handle.keys.advanced(by: index).initialize(to: key) + handle.values?.advanced(by: index).initialize(to: value) + } + + if let children = children { + for (index, child) in children.enumerated() { + handle.children!.advanced(by: index).initialize(to: child) + } + + handle.depth = handle[childAt: 0].storage.header.depth + 1 + } + + let totalChildElements = children?.reduce(0, { $0 + $1.read({ $0.subtreeCount }) }) + + handle.elementCount = keyValuePairs.count + handle.subtreeCount = keyValuePairs.count + (totalChildElements ?? 0) + } + } + + /// Converts a node to a single array. + @_spi(Testing) + public func toArray() -> [Element] { + self.read { handle in + if handle.isLeaf { + return Array((0..) { + self.element = element + self.rightChild = rightChild + } + + /// The former median element which should be propogated upward. + @usableFromInline + internal let element: Element + + /// The right product of the node split. + @usableFromInline + internal var rightChild: _Node + + /// Converts the splinter object to a node. + /// - Parameters: + /// - node: The node generating the splinter. Becomes the returned + /// node's left child. + /// - capacity: The desired capacity of the new node. + /// - Returns: A new node of `capacity` with a single element. + @inlinable + @inline(__always) + internal __consuming func toNode( + leftChild: _Node, + capacity: Int + ) -> _Node { + return _Node( + leftChild: leftChild, + seperator: element, + rightChild: rightChild, + capacity: capacity + ) + } + } +} diff --git a/Sources/SortedCollections/BTree/_Node.Storage.swift b/Sources/SortedCollections/BTree/_Node.Storage.swift new file mode 100644 index 000000000..b9cedf0af --- /dev/null +++ b/Sources/SortedCollections/BTree/_Node.Storage.swift @@ -0,0 +1,244 @@ +//===----------------------------------------------------------------------===// +// +// 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 _Node { + @usableFromInline + internal struct Header { + @inlinable + internal init( + capacity: Int, + count: Int, + subtreeCount: Int, + depth: Int, + values: UnsafeMutablePointer?, + children: UnsafeMutablePointer<_Node>? + ) { + self._internalCounts = 0 + self.values = values + self.children = children + self.subtreeCount = subtreeCount + + self.capacity = capacity + self.count = count + self.depth = depth + } + + /// Packed integer to store all node counts. + /// + /// This is represented as: + /// + /// subtreeCount + /// vvvvvvv + /// 0x0000000FFFFFFF00 + /// ^^^^^^^ ^^ + /// count depth + /// + @usableFromInline + internal var _internalCounts: UInt64 + + /// Refers to the amount of keys in the node. + @inlinable + @inline(__always) + internal var count: Int { + get { Int((_internalCounts & 0xFFFFFFF000000000) >> 40) } + set { + assert(0 <= newValue && newValue <= 0xFFFFFFF, "Invalid count.") + assert(newValue <= capacity, "Count cannot exceed capacity.") + _internalCounts &= ~0xFFFFFFF000000000 + _internalCounts |= UInt64(newValue) << 40 + } + } + + /// The total amount of keys possible to store within the node. + @inlinable + @inline(__always) + internal var capacity: Int { + get { Int((_internalCounts & 0x0000000FFFFFFF00) >> 8) } + set { + assert(0 <= newValue && newValue <= 0xFFFFFFF, "Invalid capacity.") + assert(newValue >= count, "Capacity cannot be below count.") + _internalCounts &= ~0x0000000FFFFFFF00 + _internalCounts |= UInt64(newValue) << 8 + } + } + + /// The depth of the node represented as the number of nodes below the current one. + @inlinable + @inline(__always) + internal var depth: Int { + get { Int(_internalCounts & 0x00000000000000FF) } + set { + assert(0 <= newValue && newValue <= 0xFF, "Invalid depth.") + _internalCounts &= ~0x00000000000000FF + _internalCounts |= UInt64(newValue) + } + } + + /// The total amount of elements contained underneath this node + @usableFromInline + internal var subtreeCount: Int + + /// Pointer to the buffer containing the corresponding values. + @usableFromInline + internal var values: UnsafeMutablePointer? + + /// Pointer to the buffer containing the elements. + @usableFromInline + internal var children: UnsafeMutablePointer<_Node>? + } + + /// Represents the underlying data for a node. + /// + /// Generally, this shouldn't be directly accessed. However, in performance-critical code where operating + /// on a ``_Node`` may create unwanted ARC traffic, or other concern, this provides a few low-level + /// and unsafe APIs for operations. + /// + /// A node contains a tail-allocated contiguous buffer of keys, and also may maintain pointers to buffers + /// for the corresponding values and children. + /// + /// There are two types of nodes distinguished "leaf" and "internal" nodes. Leaf nodes do not have a + /// buffer allocated for their children in the underlying storage class. + /// + /// Additionally, a node does not have a value buffer allocated in some cases. Specifically, when the + /// value type is `Void`, no value buffer is allocated. ``_Node.hasValues`` can be used to check + /// whether a value buffer exists for a given set of generic parameters. Additionally, when a value buffer + /// does not exist, ``_Node.dummyValue`` can be used to obtain a valid value within the type system + /// that can be passed to APIs. + @usableFromInline + internal class Storage: ManagedBuffer<_Node.Header, Key> { + /// Allows **read-only** access to the underlying data behind the node. + /// + /// - Parameter body: A closure with a handle which allows interacting with the node + /// - Returns: The value the closure body returns, if any. + @inlinable + @inline(__always) + internal func read(_ body: (UnsafeHandle) throws -> R) rethrows -> R { + return try self.withUnsafeMutablePointers { header, keys in + let handle = UnsafeHandle( + keys: keys, + values: header.pointee.values, + children: header.pointee.children, + header: header, + isMutable: false + ) + return try body(handle) + } + } + + /// Dangerously allows **mutable** access to the underlying data behind the node. + /// + /// This does **not** perform CoW checks and so when calling this, one must be certain that the + /// storage class is indeed unique. Generally this is the wrong function to call, and should only be used + /// when the callee has created and is guaranteed to be the only owner during the execution of the + /// update callback, _and_ it has been identified that ``_Node.update(_:)`` or other alternatives + /// result in noticable slow-down. + /// + /// - Parameter body: A closure with a handle which allows interacting with the node + /// - Returns: The value the closure body returns, if any. + /// - Warning: The underlying storage **must** be unique. + @inlinable + @inline(__always) + internal func updateGuaranteedUnique( + _ body: (UnsafeHandle) throws -> R + ) rethrows -> R { + try self.read { try body(UnsafeHandle(mutating: $0)) } + } + + /// Creates a new storage object. + /// + /// It is generally recommend to use the ``_Node.init(withCapacity:, isLeaf:)`` + /// initializer instead to create a new node. + @inlinable + @inline(__always) + internal static func create( + withCapacity capacity: Int, + isLeaf: Bool + ) -> Storage { + let storage = Storage.create(minimumCapacity: capacity) { _ in + Header( + capacity: capacity, + count: 0, + subtreeCount: 0, + depth: 0, + values: + Value.self == Void.self ? nil + : UnsafeMutablePointer.allocate(capacity: capacity), + children: isLeaf ? nil + : UnsafeMutablePointer<_Node>.allocate(capacity: capacity + 1) + ) + } + + return unsafeDowncast(storage, to: Storage.self) + } + + /// Copies an existing storage to a new storage. + /// + /// It is generally recommended to use the ``_Node.init(copyingFrom:)`` initializer. + @inlinable + @inline(__always) + internal func copy() -> Storage { + let capacity = self.header.capacity + let count = self.header.count + let subtreeCount = self.header.subtreeCount + let depth = self.header.depth + let isLeaf = self.header.children == nil + + let newStorage = Storage.create(withCapacity: capacity, isLeaf: isLeaf) + + newStorage.header.count = count + newStorage.header.subtreeCount = subtreeCount + newStorage.header.depth = depth + + self.withUnsafeMutablePointerToElements { oldKeys in + newStorage.withUnsafeMutablePointerToElements { newKeys in + newKeys.initialize(from: oldKeys, count: count) + } + } + + if _Node.hasValues { + newStorage.header.values.unsafelyUnwrapped + .initialize( + from: self.header.values.unsafelyUnwrapped, + count: count + ) + } + + newStorage.header.children? + .initialize( + from: self.header.children.unsafelyUnwrapped, + count: count + 1 + ) + + return newStorage + } + + @inlinable + deinit { + self.withUnsafeMutablePointers { header, elements in + let count = header.pointee.count + + if _Node.hasValues { + let values = header.pointee.values.unsafelyUnwrapped + values.deinitialize(count: count) + values.deallocate() + } + + + if let children = header.pointee.children { + children.deinitialize(count: count + 1) + children.deallocate() + } + + elements.deinitialize(count: header.pointee.count) + } + } + } +} diff --git a/Sources/SortedCollections/BTree/_Node.UnsafeHandle+CustomDebugStringConvertible.swift b/Sources/SortedCollections/BTree/_Node.UnsafeHandle+CustomDebugStringConvertible.swift new file mode 100644 index 000000000..d61cd01f3 --- /dev/null +++ b/Sources/SortedCollections/BTree/_Node.UnsafeHandle+CustomDebugStringConvertible.swift @@ -0,0 +1,149 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +// MARK: CustomDebugStringConvertible +extension _Node.UnsafeHandle: CustomDebugStringConvertible { + #if DEBUG + private enum PrintPosition { case start, end, middle } + private func indentDescription(_ _node: _Node.UnsafeHandle, position: PrintPosition) -> String { + let label = "(\(_node.elementCount)/\(_node.subtreeCount) \(_node.depth))" + + let spaces = String(repeating: " ", count: label.count) + + let lines = describeNode(_node).split(separator: "\n") + return lines.enumerated().map({ index, line in + var lineToInsert = line + let middle = (lines.count - 1) / 2 + if index < middle { + if position == .start { + return " " + spaces + lineToInsert + } else { + return "┃ " + spaces + lineToInsert + } + } else if index > middle { + if position == .end { + return " " + spaces + lineToInsert + } else { + return "┃ " + spaces + lineToInsert + } + } else { + switch line[line.startIndex] { + case "╺": lineToInsert.replaceSubrange(...line.startIndex, with: "━") + case "┗": lineToInsert.replaceSubrange(...line.startIndex, with: "┻") + case "┏": lineToInsert.replaceSubrange(...line.startIndex, with: "┳") + case "┣": lineToInsert.replaceSubrange(...line.startIndex, with: "╋") + default: break + } + + switch position { + case .start: return "┏━\(label)━" + lineToInsert + case .middle: return "┣━\(label)━" + lineToInsert + case .end: return "┗━\(label)━" + lineToInsert + } + } + }).joined(separator: "\n") + } + + /// A textual representation of this instance, suitable for debugging. + private func describeNode(_ _node: _Node.UnsafeHandle) -> String { + if _node.elementCount == 0 { + var result = "" + if !_node.isLeaf { + _node[childAt: 0].read { handle in + result += indentDescription(handle, position: .start) + "\n" + } + + result += "┗━ << EMPTY >>" + } else { + result = "╺━ << EMPTY >>" + } + return result + } + + var result = "" + for slot in 0..<_node.elementCount { + if !_node.isLeaf { + let child = _node[childAt: slot] + let childDescription = child.read { + indentDescription($0, position: slot == 0 ? .start : .middle) + } + result += childDescription + "\n" + } + + if _node.isLeaf { + if _node.elementCount == 1 { + result += "╺━ " + } else if slot == _node.elementCount - 1 { + result += "┗━ " + } else if slot == 0 { + result += "┏━ " + } else { + result += "┣━ " + } + } else { + result += "┣━ " + } + + if _Node.hasValues { + debugPrint(_node[keyAt: slot], terminator: ": ", to: &result) + debugPrint(_node[valueAt: slot], terminator: "", to: &result) + } else { + debugPrint(_node[keyAt: slot], terminator: "", to: &result) + } + + if !_node.isLeaf && slot == _node.elementCount - 1 { + let childDescription = _node[childAt: slot + 1].read { + indentDescription($0, position: .end) + } + result += "\n" + childDescription + } + + result += "\n" + } + return result + } + + /// A textual representation of this instance, suitable for debugging. + public var debugDescription: String { + return indentDescription(self, position: .end) + } + #else + /// A textual representation of this instance, suitable for debugging. + public var debugDescription: String { + var result = "Node<\(Key.self), \(Value.self)>([" + var first = true + for slot in 0.. _Node.Element? { + assertMutable() + + let slot = self.startSlot(forKey: key) + + if slot < self.elementCount && self[keyAt: slot] == key { + // We have found the key + if self.isLeaf { + // Deletion within a leaf + // removeElement(atSlot:) automatically adjusts node counts. + return self.removeElement(atSlot: slot) + } else { + // Deletion within an internal node + + // TODO: potentially be smarter about using the predecessor or successor. + let predecessor = self[childAt: slot].update { $0.popLastElement() } + + // Reduce the element count. + self.subtreeCount -= 1 + + // Replace the current element with the predecessor. + let element = self.exchangeElement(atSlot: slot, with: predecessor) + + // Balance the predecessor child slot, as the pop operation may have + // brought it out of balance. + self.balance(atSlot: slot) + + return element + } + } else { + if self.isLeaf { + // If we're in a leaf node and didn't find the key, it does + // not exist. + return nil + } else { + // Sanity-check + assert(slot < self.childCount, "Attempt to remove from invalid child.") + + let removedElement = self[childAt: slot].update({ + $0.removeAnyElement(forKey: key) + }) + + guard let removedElement = removedElement else { + // Could not find the key + return nil + } + + self.subtreeCount -= 1 + + // TODO: performing the remove and then balancing may result in an + // extra memmove being performed. + + // Determine if the child needs to be rebalanced + self.balance(atSlot: slot) + + return removedElement + } + } + } + + /// 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.subtreeCount, + "Cannot remove with out-of-bounds offset.") + + if self.isLeaf { + return self.removeElement(atSlot: offset) + } + + // Removing from within an internal node + var startIndex = 0 + for childSlot in 0.. _Node.Element { + assertMutable() + + if self.isLeaf { + // At a leaf, it is trivial to pop the last element + // removeElement(atSlot:) automatically updates the counts. + return self.removeElement(atSlot: 0) + } + + // Remove the subtree's element + let poppedElement = self[childAt: 0].update { $0.popFirstElement() } + + self.subtreeCount -= 1 + + self.balance(atSlot: 0) + 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 + /// ensure the tree above this is balanced. This does adjust child counts + /// + /// - Returns: The moved last element of the tree. + @inlinable + @inline(__always) + internal func popLastElement() -> _Node.Element { + assertMutable() + + if self.isLeaf { + // At a leaf, it is trivial to pop the last element + // popLastElement(at:) automatically updates the counts. + return self.removeElement(atSlot: self.elementCount - 1) + } + + // Remove the subtree's element + let poppedElement = self[childAt: self.childCount - 1].update { + $0.popLastElement() + } + + self.subtreeCount -= 1 + + self.balance(atSlot: self.childCount - 1) + return poppedElement + } +} + +// MARK: Balancing +extension _Node.UnsafeHandle { + + /// Balances a node's child at a specific slot to maintain the BTree invariants + /// and ensure none of its children underflows. + /// + /// Calling this method may make the current node underflow. Therefore, this + /// should be called upwards on the entire tree. + /// + /// If no balancing needs to occur, then this method leaves the tree unchanged. + /// + /// - Parameter slot: The slot containing the child to balance. + @inlinable + internal func balance(atSlot slot: Int) { + assertMutable() + assert(0 <= slot && slot < self.childCount, + "Cannot balance out-of-bounds slot.") + assert(!self.isLeaf, "Cannot balance leaf.") + + // No need to balance if the node is already balanced + if self[childAt: slot].read({ $0.isBalanced }) { return } + + if slot > 0 && self[childAt: slot - 1].read({ $0.isShrinkable }) { + // We can rotate from the left node to the right node + self.rotateRight(atSlot: slot - 1) + } else if slot < self.childCount - 1 && + self[childAt: slot + 1].read({ $0.isShrinkable }) { + // We can rotate from the right node to the left node + self.rotateLeft(atSlot: slot) + } else if slot == self.childCount - 1 { + // In the special case the deficient child at the end, + // it'll be merged with it's left sibling. + self.collapse(atSlot: slot - 1) + } else { + // Otherwise collapse the child with its right sibling. + self.collapse(atSlot: slot) + } + } + + /// Performs a right-rotation of the key at a given slot. + /// + /// The rotation occurs among the keys corresponding left and right child. This + /// does not perform any checks to ensure underflow does not occur. + /// + /// - Parameter slot: The slot containing the child to balance. + @inlinable + @inline(__always) + internal func rotateRight(atSlot slot: Int) { + assertMutable() + assert(0 <= slot && slot < self.elementCount, + "Cannot rotate out-of-bounds slot.") + + self[childAt: slot].update { leftChild in + assert(leftChild.elementCount > 0, "Cannot rotate from empty node.") + + // Move the old parent down to the right node + self[childAt: slot + 1].update { rightChild in + assert(rightChild.elementCount < rightChild.capacity, + "Rotating into full node.") + assert(leftChild.isLeaf == rightChild.isLeaf, + "Encountered subtrees of conflicting depth.") + + // Shift the rest of the elements right + rightChild.moveInitializeElements( + count: rightChild.elementCount, + fromSlot: 0, + toSlot: 1, of: rightChild + ) + + // Extract the parent's current element and move it down to the right + // node + let oldParentElement = self.moveElement(atSlot: slot) + if !rightChild.isLeaf { + // Move the corresponding children to the right + rightChild.moveInitializeChildren( + count: rightChild.childCount, + fromSlot: 0, + toSlot: 1, of: rightChild + ) + + // We'll extract the last child of the left node, if it exists, + // in order to cycle it. + // removeChild(atSlot:) takes care of adjusting the total element + // counts. + let newLeftChild: _Node = + leftChild.removeChild(atSlot: leftChild.childCount - 1) + + // Move the left child if applicable to the right node. + rightChild.initializeElement( + atSlot: 0, to: oldParentElement, withLeftChild: newLeftChild + ) + + rightChild.subtreeCount += newLeftChild.read({ $0.subtreeCount }) + } else { + // Move the left child if applicable to the right node. + rightChild.initializeElement(atSlot: 0, to: oldParentElement) + } + + // Move the left node's key up to the parent + // removeElement(atSlot:) takes care of adjusting the node counts. + let newParentElement = + leftChild.removeElement(atSlot: leftChild.elementCount - 1) + self.initializeElement(atSlot: slot, to: newParentElement) + + // Adjust the element counts as applicable + rightChild.elementCount += 1 + rightChild.subtreeCount += 1 + } + } + } + + /// Performs a left-rotation of the key at a given slot. + /// + /// The rotation occurs among the keys corresponding left and right child. This + /// does not perform any checks to ensure underflow does not occur. + /// + /// - Parameter slot: The slot containing the child to balance. + @inlinable + @inline(__always) + internal func rotateLeft(atSlot slot: Int) { + assertMutable() + assert(0 <= slot && slot < self.elementCount, + "Cannot rotate out-of-bounds slot.") + + self[childAt: slot].update { leftChild in + assert(leftChild.elementCount < leftChild.capacity, + "Rotating into full node.") + + // Move the old parent down to the right node + self[childAt: slot + 1].update { rightChild in + assert(rightChild.elementCount > 0, "Cannot rotate from empty node.") + assert(leftChild.isLeaf == rightChild.isLeaf, + "Encountered subtrees of conflicting depth.") + + // Move the right child if applicable to the left node. + // We'll extract the first child of the right node, if it exists, + // in order to cycle it. + // Then, cycle the parent's element down to the left child + let oldParentElement = self.moveElement(atSlot: slot) + if !leftChild.isLeaf { + // removeChild(atSlot:) takes care of adjusting the node counts. + let newRightChild = rightChild.removeChild(atSlot: 0) + + leftChild.initializeElement( + atSlot: leftChild.elementCount, + to: oldParentElement, + withRightChild: newRightChild + ) + + // Adjust the total element counts + leftChild.subtreeCount += newRightChild.read({ $0.subtreeCount }) + } else { + leftChild.initializeElement( + atSlot: leftChild.elementCount, to: oldParentElement + ) + } + + // Cycle the right child's element up to the parent. + // removeElement(atSlot:) takes care of adjusting the node counts. + let newParentElement = rightChild.removeElement(atSlot: 0) + self.initializeElement(atSlot: slot, to: newParentElement) + + leftChild.elementCount += 1 + leftChild.subtreeCount += 1 + } + } + } + + /// Collapses a slot and its children into a single child. + /// + /// This will reuse the left childs node for the new node. As a result, ensure + /// that the left node is large enough to contain both the parent and the + /// right child's contents, else this method may trap. + /// + /// As an example, calling this method on a tree such as: + /// + /// ┌─┐ + /// │5│ + /// ┌─┴─┴─┐ + /// │ │ + /// ┌─┼─┐ ┌─┼─┐ + /// │1│3│ │7│9│ + /// └─┴─┘ └─┴─┘ + /// + /// will result in a single collapsed node such as: + /// + /// ┌─┬─┬─┬─┬─┐ + /// │1│3│5│7│9│ + /// └─┴─┴─┴─┴─┘ + /// + /// - Parameter slot: The slot at which to collapse its children. + /// - Note: This method is only valid on a non-leaf nodes. + /// - Warning: Calling this may result in empty nodes and a state which breaks the + /// B-Tree invariants, ensure the tree is further balanced after this.` + @inlinable + internal func collapse(atSlot slot: Int) { + assertMutable() + assert(0 <= slot && slot < self.elementCount, + "Cannot collapse out-of-bounds slot") + assert(!self.isLeaf, "Cannot collapse a slot of a leaf node.") + + self[childAt: slot].update { leftChild in + var rightChild = self.moveChild(atSlot: slot + 1) + + // TODO: create optimized version that avoids a CoW copy for when the + // right child is shared. + rightChild.update { rightChild in + // Ensure the left handle is large enough to contain all the items + assert( + leftChild.capacity >= + leftChild.elementCount + rightChild.elementCount + 1, + "Left child undersized to contain collapsed subtree." + ) + assert(leftChild.isLeaf == rightChild.isLeaf, + "Encountered subtrees of conflicting depth.") + + // Move the remaining children + // ┌─┬─┬─┬─┐ + // │1│X│5│7│ + // └─┴─┴┬┴─┘ + // ▲ │ + // └─┘ + self.moveInitializeChildren( + count: self.elementCount - slot - 1, + fromSlot: slot + 2, + toSlot: slot + 1, of: self + ) + + // Move the element from the parent into the the left child + // + // ┌─┐ + // │5│ + // ┌──┴─┴──┐ + // │ │ │ + // ┌─┼─┐ │ ┌─┼─┐ + // │1│3│ │ │7│9│ + // └─┴─┘ │ └─┴─┘ + // ▼ + // ┌─┬─┬─┬─┬─┐ + // │1│3│5│7│9│ + // └─┴─┴─┴─┴─┘ + // The removeElement(atSlot:) takes care of adjusting the element + // count on the parent. + leftChild.initializeElement( + atSlot: leftChild.elementCount, + to: self.removeElement(atSlot: slot) + ) + + // Increment the total element count since we merely moved the + // parent element within the same subtree + self.subtreeCount += 1 + + // TODO: might be more optimal to memcpy the right child, and avoid + // potentially creating a CoW copy of it. + + // Move the right child's elements into the left child + // ┌─┐ + // │5│ + // ┌─┴─┴─┐ + // │ │ + // ┌─┼─┐ ┌─┼─┐ + // │1│3│ │7│9│ + // └─┴─┘ └─┴─┘ + // │ + // ▼ + // ┌─┬─┬─┬─┬─┐ + // │1│3│5│7│9│ + // └─┴─┴─┴─┴─┘ + rightChild.moveInitializeElements( + count: rightChild.elementCount, + fromSlot: 0, + toSlot: leftChild.elementCount + 1, of: leftChild + ) + + // Move the children of the right node to the left node + if !rightChild.isLeaf { + rightChild.moveInitializeChildren( + count: rightChild.childCount, + fromSlot: 0, + toSlot: leftChild.elementCount + 1, of: leftChild + ) + } + + // Adjust the child counts + leftChild.elementCount += rightChild.elementCount + 1 + leftChild.subtreeCount += rightChild.subtreeCount + 1 + + // Clear out the child counts for the right handle + // TODO: As part of deletion, we probably already visited the right + // node, so it's possible we just created a CoW copy of it, only do + // deallocate it. Not a big deal but an inefficiency to note. + rightChild.elementCount = 0 + rightChild.drop() + } + } + } +} diff --git a/Sources/SortedCollections/BTree/_Node.UnsafeHandle+Insertion.swift b/Sources/SortedCollections/BTree/_Node.UnsafeHandle+Insertion.swift new file mode 100644 index 000000000..dbe39b332 --- /dev/null +++ b/Sources/SortedCollections/BTree/_Node.UnsafeHandle+Insertion.swift @@ -0,0 +1,522 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +// MARK: Tree Insertions +extension _Node.UnsafeHandle { + @usableFromInline + @frozen + internal enum UpdateResult { + case updated(previousElement: _Node.Element) + case splintered(_Node.Splinter) + case inserted + + @inlinable + @inline(__always) + internal init(from splinter: _Node.Splinter?) { + if let splinter = splinter { + self = .splintered(splinter) + } else { + self = .inserted + } + } + } + + /// Insert or update an element in the tree. Starts at the current node, returning a + /// possible splinter, or the previous value for any matching key if updating a node. + /// + /// In the case of duplicates, this is marginally more efficient, however, this may update + /// any element with a key equal to the provided one in the tree. For this reason, refrain + /// from using this unless the tree is guaranteed to have unique keys, else it may have + /// inconsistent behavior. + /// + /// If a matching key is found, only the value will be updated. + /// + /// - Parameters: + /// - value: The value to insert or update. + /// - key: The key to equate. + /// - updatingKey: If the key is found, whether it should be updated. + /// - Returns: A representation of the possible results of the update/insertion. + @inlinable + @inline(__always) + internal func updateAnyValue( + _ value: Value, + forKey key: Key, + updatingKey: Bool + ) -> UpdateResult { + assertMutable() + + let insertionIndex = self.endSlot(forKey: key) + + if 0 < insertionIndex && insertionIndex <= self.elementCount && + self[keyAt: insertionIndex - 1] == key { + if updatingKey { + // TODO: Potential transient ARC traffic here. + let oldKey = self.keys.advanced(by: insertionIndex - 1).pointee + + let oldValue: Value + if _Node.hasValues { + oldValue = self.pointerToValue(atSlot: insertionIndex - 1).move() + self.pointerToValue(atSlot: insertionIndex - 1).initialize(to: value) + } else { + oldValue = _Node.dummyValue + } + + return .updated(previousElement: (oldKey, oldValue)) + } else { + let oldElement = self.exchangeElement( + atSlot: insertionIndex - 1, + with: (key, value) + ) + + return .updated(previousElement: oldElement) + } + } + + // We need to try to insert as deep as possible as first, and have the splinter + // bubble up. + if self.isLeaf { + let maybeSplinter = self.insertElement( + (key, value), + withRightChild: nil, + atSlot: insertionIndex + ) + return UpdateResult(from: maybeSplinter) + } else { + let result = self[childAt: insertionIndex].update { + $0.updateAnyValue(value, forKey: key, updatingKey: updatingKey) + } + + switch result { + case .updated: + return result + case .splintered(let splinter): + let splinter = self.insertSplinter(splinter, atSlot: insertionIndex) + return UpdateResult(from: splinter) + case .inserted: + self.subtreeCount += 1 + return .inserted + } + } + } +} + +// MARK: Immediate Node Insertions +extension _Node.UnsafeHandle { + /// Inserts a value into this node without considering the children. Be careful when using + /// this as you can violate the BTree invariants if not careful. + /// + /// - Parameters: + /// - element: The new key-value pair to insert in the node. + /// - rightChild: The new element's corresponding right-child provided iff the + /// node is not a leaf, otherwise `nil`. + /// - insertionSlot: The slot to insert the new element. + /// - Returns: A splinter object if node splintered during the insert, otherwise `nil` + /// - Warning: Ensure you insert the node in a valid order as to not break the node's + /// sorted invariant. + @inlinable + internal func insertElement( + _ element: _Node.Element, + withRightChild rightChild: _Node?, + atSlot insertionSlot: Int + ) -> _Node.Splinter? { + assertMutable() + assert(self.isLeaf == (rightChild == nil), + "A child can only be inserted iff the node is a leaf.") + + // If we have a full B-Tree, we'll need to splinter + if self.elementCount == self.capacity { + // Right median == left median for BTrees with odd capacity + let rightMedian = self.elementCount / 2 + let leftMedian = (self.elementCount - 1) / 2 + + var splinterElement: _Node.Element + var rightNode = _Node(withCapacity: self.capacity, isLeaf: self.isLeaf) + + if insertionSlot == rightMedian { + splinterElement = element + + let leftElementCount = rightMedian + let rightElementCount = self.elementCount - rightMedian + + rightNode.update { rightHandle in + self.moveInitializeElements( + count: rightElementCount, + fromSlot: rightMedian, + toSlot: 0, of: rightHandle + ) + + if !self.isLeaf { + rightHandle.children.unsafelyUnwrapped + .initialize(to: rightChild.unsafelyUnwrapped) + + self.moveInitializeChildren( + count: rightElementCount, + fromSlot: rightMedian + 1, + toSlot: 1, of: rightHandle + ) + } + + self.elementCount = leftElementCount + rightHandle.elementCount = rightElementCount + rightHandle.depth = self.depth + + self._adjustSubtreeCount(afterSplittingTo: rightHandle) + } + } else if insertionSlot > rightMedian { + // This branch is almost certainly correct + splinterElement = self.moveElement(atSlot: rightMedian) + + let insertionSlotInRightNode = insertionSlot - (rightMedian + 1) + + rightNode.update { rightHandle in + self.moveInitializeElements( + count: insertionSlotInRightNode, + fromSlot: rightMedian + 1, + toSlot: 0, of: rightHandle + ) + + self.moveInitializeElements( + count: self.elementCount - insertionSlot, + fromSlot: insertionSlot, + toSlot: insertionSlotInRightNode + 1, of: rightHandle + ) + + if !self.isLeaf { + self.moveInitializeChildren( + count: insertionSlot - rightMedian, + fromSlot: rightMedian + 1, + toSlot: 0, of: rightHandle + ) + + self.moveInitializeChildren( + count: self.elementCount - insertionSlot, + fromSlot: insertionSlot + 1, + toSlot: insertionSlotInRightNode + 2, of: rightHandle + ) + } + + rightHandle.initializeElement( + atSlot: insertionSlotInRightNode, + to: element, + withRightChild: rightChild + ) + + rightHandle.elementCount = self.elementCount - rightMedian + self.elementCount = rightMedian + rightHandle.depth = self.depth + + self._adjustSubtreeCount(afterSplittingTo: rightHandle) + } + } else { + // insertionSlot < rightMedian + splinterElement = self.moveElement(atSlot: leftMedian) + + rightNode.update { rightHandle in + self.moveInitializeElements( + count: self.elementCount - leftMedian - 1, + fromSlot: leftMedian + 1, + toSlot: 0, of : rightHandle + ) + + self.moveInitializeElements( + count: leftMedian - insertionSlot, + fromSlot: insertionSlot, + toSlot: insertionSlot + 1, of: self + ) + + if !self.isLeaf { + self.moveInitializeChildren( + count: self.elementCount - leftMedian, + fromSlot: leftMedian + 1, + toSlot: 0, of: rightHandle + ) + + self.moveInitializeChildren( + count: leftMedian - insertionSlot, + fromSlot: insertionSlot + 1, + toSlot: insertionSlot + 2, of: self + ) + } + + self.initializeElement( + atSlot: insertionSlot, + to: element, withRightChild: rightChild + ) + + rightHandle.elementCount = self.elementCount - leftMedian - 1 + self.elementCount = leftMedian + 1 + rightHandle.depth = self.depth + + self._adjustSubtreeCount(afterSplittingTo: rightHandle) + } + } + + return _Node.Splinter( + element: splinterElement, + rightChild: rightNode + ) + } else { + // TODO: potentially extract out this logic to reduce code duplication. + // Shift over elements near the insertion slot. + self.moveInitializeElements( + count: self.elementCount - insertionSlot, + fromSlot: insertionSlot, + toSlot: insertionSlot + 1, of: self + ) + + if !self.isLeaf { + self.moveInitializeChildren( + count: self.childCount - insertionSlot - 1, + fromSlot: insertionSlot + 1, + toSlot: insertionSlot + 2, of: self + ) + } + + self.initializeElement( + atSlot: insertionSlot, + to: element, withRightChild: rightChild + ) + + self.elementCount += 1 + self.subtreeCount += 1 + + return nil + } + } + + /// Inserts a splinter, attaching the children appropriately. + /// + /// This only updates the count properties by one (for the seperator). If the splinter's right child contains a + /// different amount of elements than previously existed in the tree, explicitly handle those count changes. + /// See ``_Node.UnsafeHandle._adjustSubtreeCount(afterSplittingTo:)`` to + /// potentially ease this task. + /// + /// - Parameters: + /// - splinter: The splinter object from a child + /// - insertionSlot: The slot of the child which produced the splinter + /// - Returns: Another splinter which may need to be propagated upward + @inlinable + @inline(__always) + internal func insertSplinter( + _ splinter: _Node.Splinter, + atSlot insertionSlot: Int + ) -> _Node.Splinter? { + return self.insertElement( + splinter.element, + withRightChild: splinter.rightChild, + atSlot: insertionSlot + ) + } + + /// Recomputes the total amount of elements in two nodes. + /// + /// This updates the subtree counts for both the current handle and also the provided `rightHandle`. + /// + /// Use this to recompute the tracked total element counts for the current node when it + /// splits. This performs a shallow recalculation, assuming that its children's counts are + /// already accurate. + /// + /// - Parameter rightHandle: A handle to the right-half of the split. + @inlinable + @inline(__always) + internal func _adjustSubtreeCount( + afterSplittingTo rightHandle: _Node.UnsafeHandle + ) { + assertMutable() + rightHandle.assertMutable() + + let originalTotalElements = self.subtreeCount + rightHandle.subtreeCount + var totalChildElements = 0 + + if !self.isLeaf { + // Calculate total amount of child elements + // TODO: potentially evaluate min(left.children, right.children), + // but the cost of the branch will likely exceed the cost of 1 comparison + for i in 0..= 0, + "Cannot have negative number of child elements.") + + self.subtreeCount = self.elementCount + totalChildElements + rightHandle.subtreeCount = originalTotalElements - self.subtreeCount + } + + /// Concatenates a node of the same depth to end of the current node, potentially splintering. + /// + /// This only supports a single-level of splinter, therefore + /// `node.elementCount + self.elementCount + 1` must not exceed + /// `2 * self.capacity`. + /// + /// Additionally, this **consumes** the source node which will marked to be deallocated. + /// + /// - Parameters: + /// - rightNode: A consumed node with keys greater than or equal to the seperator. + /// - seperatedBy: A seperator greater than or equal to all keys in the current node. + /// - Returns: A splinter if the node could not contain both elements. + @inlinable + internal func concatenateWith( + node rightNode: inout _Node, + seperatedBy seperator: __owned _Node.Element + ) -> _Node.Splinter? { + assertMutable() + let seperator: _Node.Element? = rightNode.update { rightHandle in + assert(self.elementCount + rightHandle.elementCount <= 2 * self.capacity, + "Parameters are too large to concatenate.") + assert(self.depth == rightHandle.depth, + "Cannot concatenate nodes of varying depths. See appendNode(_:seperatedBy:)") + + let totalElementCount = self.elementCount + rightHandle.elementCount + 1 + + // Identify if a splinter needs to occur + if totalElementCount > self.capacity { + // A splinter needs to occur + + // Split evenly (right biased). + let seperatorSlot = totalElementCount / 2 + + // Identify who needs to splinter + if seperatorSlot == self.elementCount { + // The nice case when the seperator is the splinter + return seperator + } else if seperatorSlot < self.elementCount { + // Move elements from the left node to the right node + let splinterSeperator = self.moveElement(atSlot: seperatorSlot) + + let shiftedElementCount = self.elementCount - seperatorSlot - 1 + + rightHandle.moveInitializeElements( + count: rightHandle.elementCount, + fromSlot: 0, + toSlot: shiftedElementCount + 1, + of: rightHandle + ) + + rightHandle.initializeElement( + atSlot: shiftedElementCount, + to: seperator + ) + + self.moveInitializeElements( + count: shiftedElementCount, + fromSlot: seperatorSlot + 1, + toSlot: 0, + of: rightHandle + ) + + if !self.isLeaf { + rightHandle.moveInitializeChildren( + count: rightHandle.childCount, + fromSlot: 0, + toSlot: shiftedElementCount + 1, + of: rightHandle + ) + + self.moveInitializeChildren( + count: shiftedElementCount + 1, + fromSlot: seperatorSlot + 1, + toSlot: 0, + of: rightHandle + ) + } + + // TODO: adjust counts + self.elementCount = seperatorSlot + rightHandle.elementCount = totalElementCount - seperatorSlot - 1 + + self._adjustSubtreeCount(afterSplittingTo: rightHandle) + + return splinterSeperator + } else { + // seperatorSlot > self.elementCount + // Move elements from the right node to the left node + let seperatorSlotInRightHandle = seperatorSlot - self.elementCount - 1 + let splinterSeperator = + rightHandle.moveElement(atSlot: seperatorSlotInRightHandle) + + self.initializeElement( + atSlot: self.elementCount, + to: seperator + ) + + rightHandle.moveInitializeElements( + count: seperatorSlotInRightHandle, + fromSlot: 0, + toSlot: self.elementCount + 1, + of: self + ) + + rightHandle.moveInitializeElements( + count: rightHandle.elementCount - (seperatorSlotInRightHandle + 1), + fromSlot: seperatorSlotInRightHandle + 1, + toSlot: 0, + of: rightHandle + ) + + if !self.isLeaf { + rightHandle.moveInitializeChildren( + count: seperatorSlotInRightHandle + 1, + fromSlot: 0, + toSlot: self.childCount, + of: self + ) + } + + self.elementCount = seperatorSlot + rightHandle.elementCount = totalElementCount - seperatorSlot - 1 + + self._adjustSubtreeCount(afterSplittingTo: rightHandle) + + return splinterSeperator + } + } else { + // A simple merge can be performed + self.initializeElement( + atSlot: self.elementCount, + to: seperator + ) + + rightHandle.moveInitializeElements( + count: rightHandle.elementCount, + fromSlot: 0, + toSlot: self.elementCount + 1, + of: self + ) + + if !self.isLeaf { + rightHandle.moveInitializeElements( + count: rightHandle.childCount, + fromSlot: 0, + toSlot: self.childCount, + of: self + ) + } + + self.elementCount += rightHandle.elementCount + 1 + self.subtreeCount += rightHandle.subtreeCount + 1 + + rightHandle.elementCount = 0 + rightHandle.drop() + + return nil + } + } + + // Check if it splintered + if let seperator = seperator { + return _Node.Splinter(element: seperator, rightChild: rightNode) + } else { + return nil + } + } +} diff --git a/Sources/SortedCollections/BTree/_Node.UnsafeHandle.swift b/Sources/SortedCollections/BTree/_Node.UnsafeHandle.swift new file mode 100644 index 000000000..eb46d3a49 --- /dev/null +++ b/Sources/SortedCollections/BTree/_Node.UnsafeHandle.swift @@ -0,0 +1,785 @@ +//===----------------------------------------------------------------------===// +// +// 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 _Node { + /// A handle allowing potentially mutable operations to be performed. + /// + /// An ``UnsafeHandle`` should never be constructed directly. It is instead used when modifying a + /// node through its closure APIs, such as: + /// + /// let nodeMedian = node.read { handle in + /// let medianSlot = handle.elementCount + /// return handle[elementAt: medianSlot] + /// } + /// + /// The unsafe handle provides a variety of methods to ease operations on a node. This includes + /// low-level operations such as `moveElement(_:atSlot:)` or `pointerToValue(at:)` and + /// higher level operations such as `insertElement(_:)`. + /// + /// There are two variants of an ``UnsafeHandle``. A mutable and immuable one. + /// ``_Node.update(_:)`` is an example of a method which vends a mutable unsafe handle. Only a + /// mutable unsafe handle can performed unsafe operations. This exists to ensure CoW-unsafe + /// operations are not performed. + /// + /// Whene performing operations relating to children, it is important to know that not all nodes have + /// children, or a children buffer allocated. Check the ``isLeaf`` property before accessing any + /// properties or calling any methods which may interact with children. + /// + /// Additionally, when performing operations on values, a value buffer may not always be allocated. + /// Check ``_Node.hasValues`` before performing such operations. Note that element-wise + /// operations of the handle already perform such value checks and this step is not necessary. + @usableFromInline + internal struct UnsafeHandle { + @usableFromInline + internal let header: UnsafeMutablePointer
+ + @usableFromInline + internal let keys: UnsafeMutablePointer + + @usableFromInline + internal let values: UnsafeMutablePointer? + + @usableFromInline + internal let children: UnsafeMutablePointer<_Node>? + + @inlinable + @inline(__always) + internal init( + keys: UnsafeMutablePointer, + values: UnsafeMutablePointer?, + children: UnsafeMutablePointer<_Node>?, + header: UnsafeMutablePointer
, + isMutable: Bool + ) { + self.keys = keys + self.values = values + self.children = children + self.header = header + + #if COLLECTIONS_INTERNAL_CHECKS + self._isMutable = isMutable + #endif + } + + // MARK: Mutablility Checks + #if COLLECTIONS_INTERNAL_CHECKS + @usableFromInline + internal let _isMutable: Bool + #endif + + @inlinable + @inline(__always) + internal var isMutable: Bool { + #if COLLECTIONS_INTERNAL_CHECKS + return _isMutable + #else + return true + #endif + } + + /// Check that this handle supports mutating operations. + /// Every member that mutates node data must start by calling this function. + /// This helps preventing COW violations. + /// + /// Note that this is a noop in release builds. + @inlinable + @inline(__always) + internal func assertMutable() { + #if COLLECTIONS_INTERNAL_CHECKS + assert(self._isMutable, + "Attempt to mutate a node through a read-only handle") + #endif + } + + // MARK: Invariant Checks + #if COLLECTIONS_INTERNAL_CHECKS + @inline(never) + @usableFromInline + internal func checkInvariants() { + assert(depth != 0 || self.isLeaf, + "Cannot have non-leaf of zero depth.") + } + #else + @inlinable + @inline(__always) + internal func checkInvariants() {} + #endif // COLLECTIONS_INTERNAL_CHECKS + + /// Creates a mutable version of this handle + /// - Warning: Calling this circumvents the CoW checks. + @inlinable + @inline(__always) + internal init(mutating handle: UnsafeHandle) { + self.init( + keys: handle.keys, + values: handle.values, + children: handle.children, + header: handle.header, + isMutable: true + ) + } + + // MARK: Convenience properties + @inlinable + @inline(__always) + internal var capacity: Int { header.pointee.capacity } + + /// The number of elements immediately stored in the node + @inlinable + @inline(__always) + internal var elementCount: Int { + get { header.pointee.count } + nonmutating set { assertMutable(); header.pointee.count = newValue } + } + + /// The total number of elements that this node directly or indirectly stores + @inlinable + @inline(__always) + internal var subtreeCount: Int { + get { header.pointee.subtreeCount } + nonmutating set { + assertMutable(); header.pointee.subtreeCount = newValue + } + } + + /// The depth of the node represented as the number of nodes below the current one. + @inlinable + @inline(__always) + internal var depth: Int { + get { header.pointee.depth } + nonmutating set { + assertMutable(); header.pointee.depth = newValue + } + } + + /// The number of children this node directly contains + /// - Warning: Do not access on a leaf, else will panic. + @inlinable + @inline(__always) + internal var childCount: Int { + assert(!isLeaf, "Cannot access the child count on a leaf.") + return elementCount &+ 1 + } + + /// Whether the node is the bottom-most node (a leaf) within its tree. + /// + /// This is equivalent to whether or not the node contains any keys. For leaf nodes, + /// calling certain operations which depend on children may trap. + @inlinable + @inline(__always) + internal var isLeaf: Bool { children == nil } + + /// A lower bound on the amount of keys we would want a node to contain. + /// + /// Defined as `ceil(capacity/2) - 1`. + @inlinable + @inline(__always) + internal var minimumElementCount: Int { capacity / 2 } + + /// Whether an element can be removed without triggering a rebalance. + @inlinable + @inline(__always) + internal var isShrinkable: Bool { elementCount > minimumElementCount } + + /// Whether the node contains at least the minimum number of keys. + @inlinable + @inline(__always) + internal var isBalanced: Bool { elementCount >= minimumElementCount } + + /// Whether the immediate node does not have space for an additional element + @inlinable + @inline(__always) + internal var isFull: Bool { elementCount == capacity } + + + /// Checks uniqueness of a child. + /// + /// - Warning: Will trap if executed on leaf nodes + @inlinable + @inline(__always) + internal func isChildUnique(atSlot slot: Int) -> Bool { + assert(!self.isLeaf, "Cannot access children on leaf.") + return isKnownUniquelyReferenced( + &self.pointerToChild(atSlot: slot).pointee._storage + ) + } + + /// Indicates that the node has no more children and is ready for deallocation + /// + /// It is critical to ensure that there are absolutely no children or element references + /// still owned by this node, or else it may result in a serious memory leak. + @inlinable + @inline(__always) + internal func drop() { + assertMutable() + assert(self.elementCount == 0, "Cannot drop non-empty node") + self.header.pointee.children?.deallocate() + self.header.pointee.children = nil + } + } +} + +// MARK: Subscript +extension _Node.UnsafeHandle { + @inlinable + @inline(__always) + internal subscript(elementAt slot: Int) -> _Node.Element { + get { + assert(0 <= slot && slot < self.elementCount, + "Node element subscript out of bounds.") + if _Node.hasValues { + return (key: self[keyAt: slot], value: self[valueAt: slot]) + } else { + return (key: self[keyAt: slot], value: _Node.dummyValue) + } + } + } + + + @inlinable + @inline(__always) + internal func pointerToKey( + atSlot slot: Int + ) -> UnsafeMutablePointer { + assert(0 <= slot && slot < self.elementCount, + "Node key slot out of bounds.") + return self.keys.advanced(by: slot) + } + + @inlinable + @inline(__always) + internal subscript(keyAt slot: Int) -> Key { + get { + assert(0 <= slot && slot < self.elementCount, + "Node key subscript out of bounds.") + return self.keys[slot] + } + } + + + @inlinable + @inline(__always) + internal func pointerToValue( + atSlot slot: Int + ) -> UnsafeMutablePointer { + assert(0 <= slot && slot < elementCount, + "Node value slot out of bounds.") + assert(_Node.hasValues, "Node does not have value buffer.") + return values.unsafelyUnwrapped.advanced(by: slot) + } + + @inlinable + @inline(__always) + internal subscript(valueAt slot: Int) -> Value { + get { + assert(0 <= slot && slot < self.elementCount, + "Node values subscript out of bounds.") + assert(_Node.hasValues, "Node does not have value buffer.") + return self.pointerToValue(atSlot: slot).pointee + } + + nonmutating _modify { + assertMutable() + assert(0 <= slot && slot < self.elementCount, + "Node values subscript out of bounds.") + assert(_Node.hasValues, "Node does not have value buffer.") + var value = self.pointerToValue(atSlot: slot).move() + yield &value + self.pointerToValue(atSlot: slot).initialize(to: value) + } + } + + @inlinable + @inline(__always) + internal func pointerToChild( + atSlot slot: Int + ) -> UnsafeMutablePointer<_Node> { + assert(0 <= slot && slot < self.childCount, + "Node child slot out of bounds.") + assert(!isLeaf, "Cannot access children of leaf node.") + return self.children.unsafelyUnwrapped.advanced(by: slot) + } + + /// Returns the child at a given slot as a Node object + /// - Warning: During mutations, re-accessing the same child slot is invalid. + @inlinable + @inline(__always) + internal subscript(childAt slot: Int) -> _Node { + get { + assert(!isLeaf, "Cannot access children of leaf node.") + assert(0 <= slot && slot < self.childCount, + "Node child subscript out of bounds") + return self.children.unsafelyUnwrapped[slot] + } + + nonmutating _modify { + assertMutable() + assert(!isLeaf, "Cannot modify children of leaf node.") + assert(0 <= slot && slot < self.childCount, + "Node child subscript out of bounds") + var child = self.children.unsafelyUnwrapped.advanced(by: slot).move() + defer { + self.children.unsafelyUnwrapped.advanced(by: slot).initialize(to: child) + } + yield &child + } + } +} + +// MARK: Binary Search +extension _Node.UnsafeHandle { + /// Performs O(log n) search for a key, returning the first instance when duplicates exist. This + /// returns the first possible insertion point for `key`. + /// - Parameter key: The key to search for within the node. + /// - Returns: Either the slot if the first instance of the key, otherwise + /// the valid insertion point for the key. + @inlinable + internal func startSlot(forKey key: Key) -> Int { + var start: Int = 0 + var end: Int = self.elementCount + + while end > start { + let mid = (end &- start) / 2 &+ start + + // TODO: make this info a conditional_mov + if key <= self.keys[mid] { + end = mid + } else { + start = mid &+ 1 + } + } + + return end + } + + /// Performs O(log n) search for a key, returning the last instance when duplicates exist. This + /// returns the last possible valid insertion point for `key`. + /// - Parameter key: The key to search for within the node. + /// - Returns: Either the slot after the last instance of the key, otherwise + /// the valid insertion point for the key. + @inlinable + internal func endSlot(forKey key: Key) -> Int { + var start: Int = 0 + var end: Int = self.elementCount + + while end > start { + let mid = (end &- start) / 2 &+ start + + if key >= self.keys[mid] { + start = mid &+ 1 + } else { + end = mid + } + } + + return end + } +} + +// MARK: Element-wise Buffer Operations +extension _Node.UnsafeHandle { + /// Moves elements from the current handle to a new handle. + /// + /// This moves initialized elements to an sequence of initialized element slots in the target handle. + /// + /// - Parameters: + /// - newHandle: The destination handle to write to which could be the same + /// as the source to move within a handle. + /// - sourceSlot: The offset of the source handle to move from. + /// - destinationSlot: The offset of the destination handle to write to. + /// - count: The count of values to move + /// - Warning: This does not adjust the buffer counts. + @inlinable + @inline(__always) + internal func moveInitializeElements( + count: Int, + fromSlot sourceSlot: Int, + toSlot destinationSlot: Int, + of target: _Node.UnsafeHandle + ) { + assert(sourceSlot >= 0, "Move source slot must be positive.") + assert(destinationSlot >= 0, "Move destination slot must be positive.") + assert(count >= 0, "Amount of elements to move be positive.") + assert(sourceSlot + count <= self.capacity, + "Cannot move elements beyond source buffer capacity.") + assert(destinationSlot + count <= target.capacity, + "Cannot move elements beyond destination buffer capacity.") + + self.assertMutable() + target.assertMutable() + + target.keys.advanced(by: destinationSlot) + .moveInitialize(from: self.keys.advanced(by: sourceSlot), count: count) + + if _Node.hasValues { + target.values.unsafelyUnwrapped.advanced(by: destinationSlot) + .moveInitialize( + from: self.values.unsafelyUnwrapped.advanced(by: sourceSlot), + count: count + ) + } + } + + /// Moves children from the current handle to a new handle + /// + /// This moves initialized children to an sequence of initialized element slots in the target handle. + /// + /// - Parameters: + /// - newHandle: The destination handle to write to which could be the same + /// as the source to move within a handle. + /// - sourceSlot: The offset of the source handle to move from. + /// - destinationSlot: The offset of the destintion handle to write to. + /// - count: The amount of values to move + /// - Warning: This does not adjust the buffer counts. + /// - Warning: This will trap if either the source and destination handles are leaves. + @inlinable + @inline(__always) + internal func moveInitializeChildren( + count: Int, + fromSlot sourceSlot: Int, + toSlot destinationSlot: Int, + of target: _Node.UnsafeHandle + ) { + assert(sourceSlot >= 0, "Move source slot must be positive.") + assert(destinationSlot >= 0, "Move destination slot must be positive.") + assert(count >= 0, "Amount of children to move be positive.") + assert(sourceSlot + count <= self.capacity + 1, + "Cannot move children beyond source buffer capacity.") + assert(destinationSlot + count <= target.capacity + 1, + "Cannot move children beyond destination buffer capacity.") + assert(!target.isLeaf, "Cannot move children to a leaf node") + assert(!self.isLeaf, "Cannot move children from a leaf node") + + self.assertMutable() + target.assertMutable() + + let sourcePointer = + self.children.unsafelyUnwrapped.advanced(by: sourceSlot) + + target.children.unsafelyUnwrapped.advanced(by: destinationSlot) + .moveInitialize(from: sourcePointer, count: count) + } + + /// Inserts a new element into an uninitialized slot in node. + /// + /// This ensures that a child is provided where appropriate and may trap if a right + /// child is not provided iff the node is a leaf. + /// + /// - Parameters: + /// - slot: An uninitialized slot in the buffer to insert the element into. + /// - element: The element to insert which the node will take ownership of. + /// - rightChild: The right child of the newly initialized element. Should be `nil` iff + /// node is a leaf. + /// - Warning: This does not adjust the buffer counts. + @inlinable + @inline(__always) + internal func initializeElement( + atSlot slot: Int, + to element: _Node.Element, + withRightChild rightChild: _Node? + ) { + assertMutable() + assert(0 <= slot && slot < self.capacity, + "Cannot insert beyond node capacity.") + assert(self.isLeaf == (rightChild == nil), + "A child can only be inserted iff the node is a leaf.") + + self.initializeElement(atSlot: slot, to: element) + + if let rightChild = rightChild { + self.children.unsafelyUnwrapped + .advanced(by: slot + 1) + .initialize(to: rightChild) + } + } + + /// Inserts a new element into an uninitialized slot in node. + /// + /// This ensures that a child is provided where appropriate and may trap if a left + /// child is not provided iff the node is a leaf. + /// + /// - Parameters: + /// - element: The element to insert which the node will take ownership of. + /// - slot: An uninitialized slot in the buffer to insert the element into. + /// - leftChild: The left child of the newly initialized element. Should be `nil` iff + /// node is a leaf. + /// - Warning: This does not adjust the buffer counts. + @inlinable + @inline(__always) + internal func initializeElement( + atSlot slot: Int, + to element: _Node.Element, + withLeftChild leftChild: _Node? + ) { + assertMutable() + assert(0 <= slot && slot < self.capacity, + "Cannot insert beyond node capacity.") + assert(self.isLeaf == (leftChild == nil), + "A child can only be inserted iff the node is a leaf.") + + self.initializeElement(atSlot: slot, to: element) + + if let leftChild = leftChild { + self.children.unsafelyUnwrapped + .advanced(by: slot) + .initialize(to: leftChild) + } + } + + /// Inserts a new element into an uninitialized slot in node. + /// + /// - Parameters: + /// - element: The element to insert which the node will take ownership of. + /// - slot: An uninitialized slot in the buffer to insert the element into. + /// - Warning: This does not adjust the buffer counts. + @inlinable + @inline(__always) + internal func initializeElement(atSlot slot: Int, to element: _Node.Element) { + assertMutable() + assert(0 <= slot && slot < self.capacity, + "Cannot insert beyond node capacity.") + + self.keys.advanced(by: slot).initialize(to: element.key) + if _Node.hasValues { + self.values.unsafelyUnwrapped.advanced(by: slot).initialize(to: element.value) + } + } + + /// Moves an element out of the handle and returns it. + /// + /// This may leave a hold within the node's buffer so it is critical to ensure that + /// it is filled, either by inserting a new element or some other operation. + /// + /// Additionally, this does not touch the children of the node, so it is important to + /// ensure those are correctly handled. + /// + /// - Parameter slot: The in-bounds slot of an element to move out + /// - Returns: A tuple of the key and value. + /// - Warning: This does not adjust buffer counts + @inlinable + @inline(__always) + internal func moveElement(atSlot slot: Int) -> _Node.Element { + assertMutable() + assert(0 <= slot && slot < self.elementCount, + "Attempted to move out-of-bounds element.") + + return ( + key: self.pointerToKey(atSlot: slot).move(), + value: _Node.hasValues + ? self.pointerToValue(atSlot: slot).move() + : _Node.dummyValue + ) + } + + /// Moves an element out of the handle and returns it. + /// + /// This may leave a hold within the node's buffer so it is critical to ensure that + /// it is filled, either by inserting a new child or some other operation. + /// + /// - Parameter slot: The in-bounds slot of an chile to move out + /// - Returns: The child node object. + /// - Warning: This does not adjust buffer counts + @inlinable + @inline(__always) + internal func moveChild(atSlot slot: Int) -> _Node { + assertMutable() + assert(!self.isLeaf, "Can only move a child on a non-leaf node.") + assert(0 <= slot && slot < self.childCount, + "Attempted to move out-of-bounds child.") + + return self.children.unsafelyUnwrapped.advanced(by: slot).move() + } + + /// Appends a new element. + @inlinable + @inline(__always) + internal func appendElement(_ element: _Node.Element) { + assertMutable() + assert(elementCount < capacity, "Cannot append into full node") + assert(elementCount == 0 || self[keyAt: elementCount - 1] <= element.key, + "Cannot append out-of-order element.") + + initializeElement(atSlot: elementCount, to: element) + + elementCount += 1 + subtreeCount += 1 + } + + + /// Appends a new element with a provided right child. + @inlinable + @inline(__always) + internal func appendElement( + _ element: _Node.Element, + withRightChild rightChild: _Node) { + assertMutable() + assert(!self.isLeaf, "Cannot append on leaf.") + assert(elementCount < capacity, "Cannot append into full node") + assert(elementCount == 0 || self[keyAt: elementCount - 1] <= element.key, + "Cannot append out-of-order element.") + + initializeElement( + atSlot: elementCount, + to: element, + withRightChild: rightChild + ) + + elementCount += 1 + subtreeCount += 1 + rightChild.storage.header.subtreeCount + } + + /// Swaps the element at a given slot, returning the old one. + /// - Parameters: + /// - slot: The initialized slot at which to swap the element + /// - newElement: The new element to insert + /// - Returns: The old element from the slot. + @inlinable + @inline(__always) + @discardableResult + internal func exchangeElement( + atSlot slot: Int, + with newElement: _Node.Element + ) -> _Node.Element { + assertMutable() + assert(0 <= slot && slot < self.elementCount, + "Attempted to swap out-of-bounds element.") + + let oldElement = self.moveElement(atSlot: slot) + self.initializeElement(atSlot: slot, to: newElement) + return oldElement + } + + /// Swaps the child at a given slot, returning the old one + /// - Parameters: + /// - slot: The initialized slot at which to swap the child + /// - newElement: The new child to insert + /// - Returns: The old child from the slot. + @inlinable + @inline(__always) + @discardableResult + internal func exchangeChild( + atSlot slot: Int, + with newChild: _Node + ) -> _Node { + assertMutable() + assert(!self.isLeaf, "Cannot exchange children on a leaf node.") + assert(0 <= slot && slot < self.childCount, + "Attempted to swap out-of-bounds element.") + + let oldChild = self.moveChild(atSlot: slot) + self.pointerToChild(atSlot: slot).initialize(to: newChild) + return oldChild + } + + /// Removes a child node pair from a given slot within the node. This is assumed to + /// run on a leaf. + /// + /// This is effectively the same as ``_Node.UnsafeHandle.moveChild(atSlot:)`` + /// however, this also moves surrounding elements as to prevent gaps from appearing + /// within the buffer. + /// + /// This does not touch the elements in the node, so it is important to ensure those are + /// correctly handled. + /// + /// This operation does adjust the stored subtree count on the node as appropriate. + /// + /// - Parameter slot: The slot to remove which must be in-bounds. + /// - Returns: The child that was removed. + /// - Warning: This performs neither balancing, rotation, or count updates. + @inlinable + @inline(__always) + internal func removeChild(atSlot slot: Int) -> _Node { + assertMutable() + assert(0 <= slot && slot < self.childCount, + "Attempt to remove out-of-bounds child.") + + let child = self.moveChild(atSlot: slot) + + // Shift everything else over to the left + self.moveInitializeChildren( + count: self.childCount - slot - 1, + fromSlot: slot + 1, + toSlot: slot, of: self + ) + + self.subtreeCount -= child.read({ $0.subtreeCount }) + + return child + } + + /// Removes a key-value pair from a given slot within the node. + /// + /// This does not touch the children of the node, so it is important to ensure those are + /// correctly handled. + /// + /// This operation adjusts the stored counts on the node as appropriate. + /// + /// - Parameter slot: The slot to remove which must be in-bounds. + /// - Returns: The element that was removed. + /// - Warning: This does not perform any balancing or rotation. + @inlinable + @inline(__always) + @discardableResult + internal func removeElement(atSlot slot: Int) -> _Node.Element { + assertMutable() + assert(0 <= slot && slot < self.elementCount, + "Attempt to remove out-of-bounds element.") + + let element = self.moveElement(atSlot: slot) + + // Shift everything else over to the left + self.moveInitializeElements( + count: self.elementCount - slot - 1, + fromSlot: slot + 1, + toSlot: slot, of: self + ) + + self.elementCount -= 1 + self.subtreeCount -= 1 + + return element + } + + /// Removes a **key** from a given slot within the node. + /// + /// This does not touch the children of the node, so it is important to ensure those are + /// correctly handled. + /// + /// This assumes the value buffer at the slot is deallocated + /// + /// This operation adjusts the stored counts on the node as appropriate. + /// + /// - Parameter slot: The slot to remove which must be in-bounds. + /// - Returns: The element that was removed. + /// - Warning: This does not perform any balancing or rotation. + @inlinable + @inline(__always) + @discardableResult + internal func removeElementWithoutValue(atSlot slot: Int) -> Key { + assertMutable() + assert(0 <= slot && slot < self.elementCount, + "Attempt to remove out-of-bounds element.") + + let key = self.pointerToKey(atSlot: slot).move() + + // Shift everything else over to the left + self.moveInitializeElements( + count: self.elementCount - slot - 1, + fromSlot: slot + 1, + toSlot: slot, of: self + ) + + self.elementCount -= 1 + self.subtreeCount -= 1 + + return key + } +} diff --git a/Sources/SortedCollections/BTree/_Node.swift b/Sources/SortedCollections/BTree/_Node.swift new file mode 100644 index 000000000..3fca25433 --- /dev/null +++ b/Sources/SortedCollections/BTree/_Node.swift @@ -0,0 +1,362 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +/// A single node within a B-Tree, containing keys, values, and children. +/// +/// A node is merely struct wrapper of the ``_Node.Storage`` class. This does not and should not +/// contain any other properties. By using a ``_Node`` over the underlying storage class, operations can +/// be performed with automatic copy-on-write (CoW) checks such as through the ``update(_:)`` +/// method. +/// +/// Refer to ``_Node.Storage`` for more information on the allocation and structure of the underlying +/// buffers of a node. +/// +/// You cannot operate or read directly from a node. Instead use `read(_:)` and `update(_:)` to make +/// modifications to a node. +/// +/// let nodeMedian = node.read { handle in +/// let medianSlot = handle.elementCount +/// return handle[elementAt: medianSlot] +/// } +/// +/// Refer to ``_Node.UnsafeHandle`` for the APIs available when operating on a node in such a +/// manner. +@usableFromInline +internal struct _Node { + @usableFromInline + typealias Element = (key: Key, value: Value) + + /// An optional parameter to the storage. Use ``storage`` instead + /// + /// This will never be `nil` during a valid access of ``_Node``. However, to support moving the + /// underlying storage instance for internal and unsafe operations, this is made optional as an + /// implementation artifact. + @usableFromInline + internal var _storage: Storage? + + /// Strong reference to the node's underlying data + @inlinable + @inline(__always) + internal var storage: Storage { _storage.unsafelyUnwrapped } + + /// An instance of a ``_Node`` with no underlying storage allocated. + /// + /// Use this when you need a dummy node. It is invalid to ever attempt to read or write to this node. + @inlinable + @inline(__always) + internal static var dummy: _Node { + _Node(_underlyingStorage: nil) + } + + /// Creates a node with a potentially empty underlying storage. + @inlinable + @inline(__always) + internal init(_underlyingStorage storage: Storage?) { + self._storage = storage + } + + /// Creates a node wrapping a Storage object in order to interact with it. + /// - Parameter storage: Underlying node storage. + @inlinable + @inline(__always) + internal init(_ storage: Storage) { + self._storage = storage + } + + /// Creates a new node from a left, right, and seperator. + @inlinable + internal init( + leftChild: __owned _Node, + seperator: __owned Element, + rightChild: __owned _Node, + capacity: Int + ) { + assert( + leftChild.storage.header.depth == rightChild.storage.header.depth, + "Left and right nodes of a splinter must have equal depth" + ) + + self.init(withCapacity: capacity, isLeaf: false) + self.storage.updateGuaranteedUnique { handle in + handle.keys.initialize(to: seperator.key) + if _Node.hasValues { + handle.values.unsafelyUnwrapped.initialize(to: seperator.value) + } + + handle.children.unsafelyUnwrapped.initialize(to: leftChild) + handle.children.unsafelyUnwrapped.advanced(by: 1).initialize(to: rightChild) + + handle.elementCount = 1 + handle.subtreeCount = 1 + + leftChild.storage.header.subtreeCount + + rightChild.storage.header.subtreeCount + handle.depth = leftChild.storage.header.depth + 1 + } + } + + /// Creates a new node with values modified by a transformation closure + @inlinable + @inline(__always) + internal init( + mappingFrom existingNode: _Node, + _ transform: (T) throws -> Value + ) rethrows { + let isLeaf = existingNode.storage.header.children == nil + _storage = .create( + withCapacity: existingNode.storage.header.capacity, + isLeaf: isLeaf + ) + + let elementCount = existingNode.storage.header.count + storage.header.count = elementCount + storage.header.subtreeCount = existingNode.storage.header.subtreeCount + + storage.withUnsafeMutablePointerToElements { keys in + existingNode.storage.withUnsafeMutablePointerToElements { sourceKeys in + keys.initialize(from: sourceKeys, count: existingNode.storage.header.count) + } + } + + for i in 0...dummyValue + ) + + storage.header.values.unsafelyUnwrapped + .advanced(by: i) + .initialize(to: newValue) + } + + if !isLeaf { + let oldChild = existingNode.storage.header + .children.unsafelyUnwrapped + .advanced(by: i) + .pointee + + storage.header.children.unsafelyUnwrapped + .advanced(by: i) + .initialize(to: try _Node(mappingFrom: oldChild, transform)) + } + } + + if !isLeaf { + let oldChild = existingNode.storage.header + .children.unsafelyUnwrapped + .advanced(by: elementCount) + .pointee + + storage.header.children.unsafelyUnwrapped + .advanced(by: elementCount) + .initialize(to: try _Node(mappingFrom: oldChild, transform)) + } + } + + /// Creates a node which copies the storage of an existing node. + @inlinable + @inline(__always) + internal init(copyingFrom oldNode: _Node) { + self._storage = oldNode.storage.copy() + } + + /// Creates and allocates a new node. + /// - Parameters: + /// - capacity: the maximum potential size of the node. + /// - isLeaf: whether or not the node is a leaf (it does not have any children). + @inlinable + internal init(withCapacity capacity: Int, isLeaf: Bool) { + self._storage = Storage.create(withCapacity: capacity, isLeaf: isLeaf) + } + + /// Whether the B-Tree has values associated with the keys + @inlinable + @inline(__always) + internal static var hasValues: Bool { _fastPath(Value.self != Void.self) } + + /// Dummy value for when a B-Tree does not maintain a value buffer. + /// + /// - Warning: Traps when used on a tree with a value buffer. + @inlinable + @inline(__always) + internal static var dummyValue: Value { + assert(!hasValues, "Cannot get dummy value on tree with value buffer.") + return unsafeBitCast((), to: Value.self) + } +} + +// MARK: Join Subroutine +extension _Node { + /// Joins the current node with another node of potentially differing depths. + /// + /// If you know that your nodes are the same depth, then use + /// ``_Node.UnsafeHandle.concatenateWith(node:seperatedBy:)``. + /// + /// - Parameters: + /// - leftNode:A well-formed node with elements less than or equal to `seperator`. This + /// node is **consumed and invalided** when this method is called. + /// - rightNode: A well-formed node with elements greater than or equal to `seperator`. This + /// node is **consumed and invalided** when this method is called. + /// - seperator: An element greater than or equal to all elements in the current node. + /// - Returns: A new node containing both the right and left node combined. This may or may not be + /// referentially identical to one of the old nodes. + @inlinable + internal static func join( + _ leftNode: inout _Node, + with rightNode: inout _Node, + seperatedBy seperator: __owned _Node.Element, + capacity: Int + ) -> _Node { + let leftNodeDepth = leftNode.storage.header.depth + let leftNodeSubtreeCount = leftNode.storage.header.subtreeCount + let rightNodeDepth = rightNode.storage.header.depth + let rightNodeSubtreeCount = rightNode.storage.header.subtreeCount + + func prepending( + atDepth depth: Int, + onto node: inout _Node + ) -> _Node.Splinter? { + if depth == 0 { + let splinter = leftNode.update { + $0.concatenateWith(node: &node, seperatedBy: seperator) + } + node = leftNode + return splinter + } else { + return node.update { handle in + let splinter = prepending(atDepth: depth - 1, onto: &handle[childAt: 0]) + handle.subtreeCount += leftNodeSubtreeCount + if let splinter = splinter { + return handle.insertSplinter(splinter, atSlot: 0) + } else { + handle.subtreeCount += 1 + return nil + } + } + } + } + + func appending( + atDepth depth: Int, + onto node: _Node.UnsafeHandle + ) -> _Node.Splinter? { + assert(node.depth >= depth, "Cannot graft at a depth deeper than the node.") + + if depth == 0 { + // Graft at the current node + return node.concatenateWith(node: &rightNode, seperatedBy: seperator) + } else { + let endSlot = node.childCount - 1 + let splinter = node[childAt: endSlot].update { + appending(atDepth: depth - 1, onto: $0) + } + + node.subtreeCount += rightNodeSubtreeCount + + if let splinter = splinter { + return node.insertSplinter(splinter, atSlot: endSlot) + } else { + node.subtreeCount += 1 + return nil + } + } + } + + if leftNodeDepth >= rightNodeDepth { + let splinter = leftNode.update { + appending(atDepth: leftNodeDepth - rightNodeDepth, onto: $0) + } + if let splinter = splinter { + return splinter.toNode(leftChild: leftNode, capacity: capacity) + } else { + return leftNode + } + } else { + let splinter = prepending(atDepth: rightNodeDepth - leftNodeDepth, onto: &rightNode) + if let splinter = splinter { + return splinter.toNode(leftChild: rightNode, capacity: capacity) + } else { + return rightNode + } + } + } +} + +// MARK: CoW +extension _Node { + /// Allows **read-only** access to the underlying data behind the node. + /// + /// - Parameter body: A closure with a handle which allows interacting with the node + /// - Returns: The value the closure body returns, if any. + @inlinable + @inline(__always) + internal func read(_ body: (UnsafeHandle) throws -> R) rethrows -> R { + return try self.storage.read(body) + } + + /// Allows mutable access to the underlying data behind the node. + /// + /// - Parameter body: A closure with a handle which allows interacting with the node + /// - Returns: The value the closure body returns, if any. + @inlinable + @inline(__always) + internal mutating func update( + _ body: (UnsafeHandle) throws -> R + ) rethrows -> R { + self.ensureUnique() + return try self.read { handle in + defer { handle.checkInvariants() } + return try body(UnsafeHandle(mutating: handle)) + } + } + + /// Allows mutable access to the underlying data behind the node. + /// - Parameters: + /// - isUnique: Whether the node is unique or needs to be copied + /// - body: A closure with a handle which allows interacting with the node + /// - Returns: The value the closure body returns, if any. + @inlinable + @inline(__always) + internal mutating func update( + isUnique: Bool, + _ body: (UnsafeHandle) throws -> R + ) rethrows -> R { + if !isUnique { + self = _Node(copyingFrom: self) + } + + return try self.read { handle in + defer { handle.checkInvariants() } + return try body(UnsafeHandle(mutating: handle)) + } + } + + /// Ensure that this storage refers to a uniquely held buffer by copying + /// elements if necessary. + @inlinable + internal mutating func ensureUnique() { + if !isKnownUniquelyReferenced(&self._storage) { + self = _Node(copyingFrom: self) + } + } +} + +// MARK: Equatable +extension _Node: Equatable { + /// Whether two nodes are the same underlying reference in memory. + /// - Warning: This **does not** compare the keys at all. + @inlinable + @inline(__always) + internal static func ==(lhs: _Node, rhs: _Node) -> Bool { + return lhs.storage === rhs.storage + } +} + diff --git a/Sources/SortedCollections/SortedDictionary/SortedDictionary+BidirectionalCollection.swift b/Sources/SortedCollections/SortedDictionary/SortedDictionary+BidirectionalCollection.swift new file mode 100644 index 000000000..a9027a3fa --- /dev/null +++ b/Sources/SortedCollections/SortedDictionary/SortedDictionary+BidirectionalCollection.swift @@ -0,0 +1,175 @@ +//===----------------------------------------------------------------------===// +// +// 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: BidirectionalCollection { + /// The number of elements in the sorted dictionary. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var count: Int { self._root.count } + + /// A Boolean value that indicates whether the dictionary is empty. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var isEmpty: Bool { self._root.isEmpty } + + /// The position of the first element in a nonempty dictionary. + /// + /// If the collection is empty, `startIndex` is equal to `endIndex`. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public var startIndex: Index { Index(self._root.startIndex) } + + /// The dictionary's "past the end" position---that is, the position one + /// greater than the last valid subscript argument. + /// + /// If the collection is empty, `endIndex` is equal to `startIndex`. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var endIndex: Index { Index(self._root.endIndex) } + + /// Returns the distance between two indices. + /// + /// - 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. + /// - Complexity: O(1) + @inlinable + public func distance(from start: Index, to end: Index) -> Int { + start._index.ensureValid(forTree: self._root) + end._index.ensureValid(forTree: self._root) + return self._root.distance(from: start._index, to: end._index) + } + + /// Replaces the given index with its successor. + /// + /// - Parameter index: A valid index of the collection. `index` must be less than `endIndex`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + public func formIndex(after index: inout Index) { + index._index.ensureValid(forTree: self._root) + self._root.formIndex(after: &index._index) + } + + /// Returns the position immediately after the given index. + /// + /// - Parameter index: A valid index of the collection. `index` must be less than `endIndex`. + /// - Returns: The index value immediately after `index`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + public func index(after index: Index) -> Index { + index._index.ensureValid(forTree: self._root) + return Index(self._root.index(after: index._index)) + } + + /// Replaces the given index with its predecessor. + /// + /// - Parameter index: A valid index of the collection. `index` must be greater + /// than `startIndex`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + public func formIndex(before index: inout Index) { + index._index.ensureValid(forTree: self._root) + self._root.formIndex(before: &index._index) + } + + /// Returns the position immediately before the given index. + /// + /// - Parameter index: A valid index of the collection. `index` must be greater + /// than `startIndex`. + /// - Returns: The index value immediately before `index`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + public func index(before index: Index) -> Index { + index._index.ensureValid(forTree: self._root) + return Index(self._root.index(before: index._index)) + } + + /// Offsets the given index by the specified distance. + /// + /// The value passed as distance must not offset i beyond the bounds of the collection. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + public func formIndex(_ i: inout Index, offsetBy distance: Int) { + i._index.ensureValid(forTree: self._root) + self._root.formIndex(&i._index, offsetBy: distance) + } + + /// Returns an index that is the specified distance from the given index. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - Returns: An index offset by `distance` from the index `i`. If `distance` + /// is positive, this is the same value as the result of `distance` calls to + /// `index(after:)`. If `distance` is negative, this is the same value as the + /// result of `abs(distance)` calls to `index(before:)`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + public func index(_ i: Index, offsetBy distance: Int) -> Index { + i._index.ensureValid(forTree: self._root) + return Index(self._root.index(i._index, offsetBy: distance)) + } + + /// Returns an index that is the specified distance from the given index, unless that distance is beyond + /// a given limiting index. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - limit: A valid index of the collection to use as a limit. If `distance > 0`, a limit that is less + /// than `i` has no effect. Likewise, if `distance < 0`, a limit that is greater than `i` has + /// no effect. + /// - Returns: An index offset by `distance` from the index `i`, unless that index would be + /// beyond `limit` in the direction of movement. In that case, the method returns `nil`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + public func index(_ i: Index, offsetBy distance: Int, limitedBy limit: Index) -> Index? { + i._index.ensureValid(forTree: self._root) + limit._index.ensureValid(forTree: self._root) + if let i = self._root.index(i._index, offsetBy: distance, limitedBy: limit._index) { + return Index(i) + } else { + return nil + } + } + + /// Offsets the given index by the specified distance, or so that it equals the given limiting index. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - limit: A valid index of the collection to use as a limit. If `distance > 0`, a limit that is less + /// than `i` has no effect. Likewise, if `distance < 0`, a limit that is greater than `i` has + /// no effect. + /// - Returns: `true` if `i` has been offset by exactly `distance` steps without going beyond + /// `limit`; otherwise, `false`. When the return value is `false`, the value of `i` is + /// equal to `limit`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + internal func formIndex(_ i: inout Index, offsetBy distance: Int, limitedBy limit: Self.Index) -> Bool { + i._index.ensureValid(forTree: self._root) + limit._index.ensureValid(forTree: self._root) + return self._root.formIndex(&i._index, offsetBy: distance, limitedBy: limit._index) + } +} diff --git a/Sources/SortedCollections/SortedDictionary/SortedDictionary+Codable.swift b/Sources/SortedCollections/SortedDictionary/SortedDictionary+Codable.swift new file mode 100644 index 000000000..496be575a --- /dev/null +++ b/Sources/SortedCollections/SortedDictionary/SortedDictionary+Codable.swift @@ -0,0 +1,75 @@ +//===----------------------------------------------------------------------===// +// +// 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() + try self.forEach { (key, value) in + 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() + var builder = _Tree.Builder(deduplicating: true) + var previousKey: Key? = nil + + 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) + + guard previousKey == nil || previousKey! < key else { + let context = DecodingError.Context( + codingPath: container.codingPath, + debugDescription: "Decoded elements out of order.") + throw DecodingError.dataCorrupted(context) + } + + builder.append((key, value)) + previousKey = key + } + + self.init(_rootedAt: builder.finish()) + } +} diff --git a/Sources/SortedCollections/SortedDictionary/SortedDictionary+CustomReflectable.swift b/Sources/SortedCollections/SortedDictionary/SortedDictionary+CustomReflectable.swift new file mode 100644 index 000000000..7f34e156a --- /dev/null +++ b/Sources/SortedCollections/SortedDictionary/SortedDictionary+CustomReflectable.swift @@ -0,0 +1,17 @@ +//===----------------------------------------------------------------------===// +// +// 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: CustomReflectable { + /// The custom mirror for this instance. + public var customMirror: Mirror { + Mirror(self, unlabeledChildren: self, displayStyle: .dictionary) + } +} diff --git a/Sources/SortedCollections/SortedDictionary/SortedDictionary+CustomStringConvertible.swift b/Sources/SortedCollections/SortedDictionary/SortedDictionary+CustomStringConvertible.swift new file mode 100644 index 000000000..49b63d2e4 --- /dev/null +++ b/Sources/SortedCollections/SortedDictionary/SortedDictionary+CustomStringConvertible.swift @@ -0,0 +1,52 @@ +//===----------------------------------------------------------------------===// +// +// 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: CustomStringConvertible, CustomDebugStringConvertible { + @inlinable + public var description: String { + if isEmpty { return "[:]" } + var result = "[" + var first = true + for (key, value) in self { + if first { + first = false + } else { + result += ", " + } + result += "\(key): \(value)" + } + result += "]" + return result + } + + @inlinable + public var debugDescription: String { + var result = "SortedDictionary<\(Key.self), \(Value.self)>(" + if isEmpty { + result += "[:]" + } else { + result += "[" + var first = true + for (key, value) in self { + if first { + first = false + } else { + result += ", " + } + + debugPrint(key, value, separator: ": ", terminator: "", to: &result) + } + result += "]" + } + result += ")" + return result + } +} diff --git a/Sources/SortedCollections/SortedDictionary/SortedDictionary+Equatable.swift b/Sources/SortedCollections/SortedDictionary/SortedDictionary+Equatable.swift new file mode 100644 index 000000000..cec7ae101 --- /dev/null +++ b/Sources/SortedCollections/SortedDictionary/SortedDictionary+Equatable.swift @@ -0,0 +1,33 @@ +//===----------------------------------------------------------------------===// +// +// 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: Equatable where Value: Equatable { + /// Returns a Boolean value indicating whether two values are equal. + /// + /// Equality is the inverse of inequality. For any values `a` and `b`, + /// `a == b` implies that `a != b` is false. + /// + /// - Parameters: + /// - lhs: A value to compare. + /// - rhs: Another value to compare. + /// - Complexity: O(`self.count`) + @inlinable + public static func ==(lhs: SortedDictionary, rhs: SortedDictionary) -> Bool { + // TODO: optimize/benchmarking by comparing node identity/shared subtrees. + if lhs.count != rhs.count { return false } + for ((k1, v1), (k2, v2)) in zip(lhs, rhs) { + if k1 != k2 || v1 != v2 { + return false + } + } + return true + } +} diff --git a/Sources/SortedCollections/SortedDictionary/SortedDictionary+ExpressibleByDictionaryLiteral.swift b/Sources/SortedCollections/SortedDictionary/SortedDictionary+ExpressibleByDictionaryLiteral.swift new file mode 100644 index 000000000..289b1761c --- /dev/null +++ b/Sources/SortedCollections/SortedDictionary/SortedDictionary+ExpressibleByDictionaryLiteral.swift @@ -0,0 +1,33 @@ +//===----------------------------------------------------------------------===// +// +// 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: ExpressibleByDictionaryLiteral { + /// Creates a new sorted dictionary from the contents of a dictionary + /// literal. + /// + /// Duplicate elements in the literal are allowed, but the resulting + /// set will only contain the first occurrence of each. + /// + /// Do not call this initializer directly. It is used by the compiler when you + /// use a dictionary literal. Instead, create a new ordered dictionary using a + /// dictionary literal as its value by enclosing a comma-separated list of + /// values in square brackets. You can use an array literal anywhere a set is + /// expected by the type context. + /// + /// - Parameter elements: A variadic list of key-value pairs for the new + /// dictionary. + /// + /// - Complexity: O(`n log n`) + @inlinable + public init(dictionaryLiteral elements: (Key, Value)...) { + self.init(keysWithValues: elements) + } +} diff --git a/Sources/SortedCollections/SortedDictionary/SortedDictionary+Hashable.swift b/Sources/SortedCollections/SortedDictionary/SortedDictionary+Hashable.swift new file mode 100644 index 000000000..7a6660913 --- /dev/null +++ b/Sources/SortedCollections/SortedDictionary/SortedDictionary+Hashable.swift @@ -0,0 +1,26 @@ +//===----------------------------------------------------------------------===// +// +// 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: Hashable where Key: Hashable, Value: Hashable { + /// Hashes the essential components of this value by feeding them + /// into the given hasher. + /// - Parameter hasher: The hasher to use when combining + /// the components of this instance. + /// - Complexity: O(`self.count`) + @inlinable + public func hash(into hasher: inout Hasher) { + hasher.combine(self.count) + for (key, value) in self { + hasher.combine(key) + hasher.combine(value) + } + } +} diff --git a/Sources/SortedCollections/SortedDictionary/SortedDictionary+Initializers.swift b/Sources/SortedCollections/SortedDictionary/SortedDictionary+Initializers.swift new file mode 100644 index 000000000..6916a6b46 --- /dev/null +++ b/Sources/SortedCollections/SortedDictionary/SortedDictionary+Initializers.swift @@ -0,0 +1,144 @@ +//===----------------------------------------------------------------------===// +// +// 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 { + /// Creates a dictionary from a sequence of key-value pairs. + /// + /// If duplicates are encountered the last instance of the key-value pair is the one + /// that is kept. + /// + /// - Parameter keysAndValues: A sequence of key-value pairs to use + /// for the new dictionary. + /// - Complexity: O(`n` * log(`n`)) where `n` is the number of elements in + /// the sequence. + @inlinable + @inline(__always) + public init( + keysWithValues keysAndValues: S + ) where S: Sequence, S.Element == (key: Key, value: Value) { + self.init() + + for (key, value) in keysAndValues { + self._root.updateAnyValue(value, forKey: key, updatingKey: true) + } + } + + /// Creates a dictionary from a sequence of key-value pairs. + /// + /// If duplicates are encountered the last instance of the key-value pair is the one + /// that is kept. + /// + /// - Parameter keysAndValues: A sequence of key-value pairs to use + /// for the new dictionary. + /// - Complexity: O(`n` * log(`n`)) where `n` is the number of elements in + /// the sequence. + @inlinable + @inline(__always) + public init( + keysWithValues keysAndValues: S + ) where S: Sequence, S.Element == (Key, Value) { + self.init() + + for (key, value) in keysAndValues { + self._root.updateAnyValue(value, forKey: key) + } + } + + /// Creates a dictionary from a sequence of **sorted** key-value pairs. + /// + /// This is a more efficient alternative to ``init(keysWithValues:)`` which offers + /// better asymptotic performance, and also reduces memory usage when constructing a + /// sorted dictionary on a pre-sorted sequence. + /// + /// - Parameter keysAndValues: A sequence of key-value pairs in non-decreasing + /// comparison order for the new dictionary. + /// - Complexity: O(`n`) where `n` is the number of elements in the + /// sequence. + @inlinable + @inline(__always) + public init( + sortedKeysWithValues keysAndValues: S + ) where S: Sequence, S.Element == (key: Key, value: Value) { + var builder = _Tree.Builder() + + var previousKey: Key? = nil + for (key, value) in keysAndValues { + precondition(previousKey == nil || previousKey! < key, + "Sequence out of order.") + builder.append((key, value)) + previousKey = key + } + + self.init(_rootedAt: builder.finish()) + } + + /// Creates a dictionary from a sequence of **sorted** key-value pairs. + /// + /// This is a more efficient alternative to ``init(keysWithValues:)`` which offers + /// better asymptotic performance, and also reduces memory usage when constructing a + /// sorted dictionary on a pre-sorted sequence. + /// + /// - Parameter keysAndValues: A sequence of key-value pairs in non-decreasing + /// comparison order for the new dictionary. + /// - Complexity: O(`n`) where `n` is the number of elements in the + /// sequence. + @inlinable + @inline(__always) + public init( + sortedKeysWithValues keysAndValues: S + ) where S: Sequence, S.Element == (Key, Value) { + var builder = _Tree.Builder() + + var previousKey: Key? = nil + for (key, value) in keysAndValues { + precondition(previousKey == nil || previousKey! < key, + "Sequence out of order.") + builder.append((key, value)) + previousKey = key + } + + self.init(_rootedAt: builder.finish()) + } + + /// Creates a new sorted dictionary whose keys are the groupings returned + /// by the given closure and whose values are arrays of the elements that + /// returned each key. + /// + /// The arrays in the “values” position of the new sorted dictionary each contain at least + /// one element, with the elements in the same order as the source sequence. + /// The following example declares an array of names, and then creates a sorted dictionary + /// from that array by grouping the names by first letter: + /// + /// let students = ["Kofi", "Abena", "Efua", "Kweku", "Akosua"] + /// let studentsByLetter = SortedDictionary(grouping: students, by: { $0.first! }) + /// // ["A": ["Abena", "Akosua"], "E": ["Efua"], "K": ["Kofi", "Kweku"]] + /// + /// The new `studentsByLetter` sorted dictionary has three entries, with students’ names + /// grouped by the keys `"A"`, `"E"`, and `"K"` + /// + /// + /// - Parameters: + /// - values: A sequence of values to group into a dictionary. + /// - keyForValue: A closure that returns a key for each element in values. + @inlinable + public init( + grouping values: S, + by keyForValue: (S.Element) throws -> Key + ) rethrows where Value == [S.Element], S : Sequence { + self.init() + for value in values { + let key = try keyForValue(value) + self.modifyValue(forKey: key, default: []) { group in + group.append(value) + } + } + } +} diff --git a/Sources/SortedCollections/SortedDictionary/SortedDictionary+Keys.swift b/Sources/SortedCollections/SortedDictionary/SortedDictionary+Keys.swift new file mode 100644 index 000000000..4c818ba90 --- /dev/null +++ b/Sources/SortedCollections/SortedDictionary/SortedDictionary+Keys.swift @@ -0,0 +1,253 @@ +//===----------------------------------------------------------------------===// +// +// 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 { + /// A view of an sorted dictionary's Keys as a standalone collection. + public struct Keys { + @usableFromInline + internal var _base: SortedDictionary + + @inlinable + @inline(__always) + internal init(_base: SortedDictionary) { + self._base = _base + } + } +} + +extension SortedDictionary.Keys: Sequence { + /// The element type of the collection. + public typealias Element = Key + + /// The type that allows iteration over the collection's elements. + public struct Iterator: IteratorProtocol { + @usableFromInline + internal var _iterator: SortedDictionary.Iterator + + @inlinable + @inline(__always) + internal init(_ _iterator: SortedDictionary.Iterator) { + self._iterator = _iterator + } + + @inlinable + @inline(__always) + public mutating func next() -> Element? { + _iterator.next()?.key + } + } + + @inlinable + @inline(__always) + public __consuming func makeIterator() -> Iterator { + Iterator(_base.makeIterator()) + } +} + +extension SortedDictionary.Keys: BidirectionalCollection { + /// The index type for a dictionary's keys view. + public typealias Index = SortedDictionary.Index + + /// The number of elements in the collection. + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var count: Int { self._base.count } + + /// A Boolean value that indicates whether the collection is empty. + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var isEmpty: Bool { self._base.isEmpty } + + /// The position of the first element in a nonempty collection. + /// + /// If the collection is empty, `startIndex` is equal to `endIndex`. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public var startIndex: Index { self._base.startIndex } + + /// The collection's "past the end" position---that is, the position one + /// greater than the last valid subscript argument. + /// + /// If the collection is empty, `endIndex` is equal to `startIndex`. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var endIndex: Index { self._base.endIndex } + + /// Returns the distance between two indices. + /// + /// - 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`. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public func distance(from start: Index, to end: Index) -> Int { + self._base.distance(from: start, to: end) + } + + /// Replaces the given index with its successor. + /// + /// The specified index must be a valid index less than `endIndex`, or the + /// returned value won't be a valid index in the collection. + /// + /// - Parameter i: A valid index of the collection. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public func formIndex(after index: inout Index) { + self._base.formIndex(after: &index) + } + + /// Returns the position immediately after the given index. + /// + /// The specified index must be a valid index less than `endIndex`, or the + /// returned value won't be a valid index in the collection. + /// + /// - Parameter i: A valid index of the collection. + /// + /// - Returns: The index immediately after `i`. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public func index(after index: Index) -> Index { + self._base.index(after: index) + } + + /// Replaces the given index with its successor. + /// + /// The specified index must be a valid index less than `endIndex`, or the + /// returned value won't be a valid index in the collection. + /// + /// - Parameter i: A valid index of the collection. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public func formIndex(before index: inout Index) { + self._base.formIndex(before: &index) + } + + /// Returns the position immediately before the given index. + /// + /// The specified index must be a valid index greater than `startIndex`, or + /// the returned value won't be a valid index in the collection. + /// + /// - Parameter i: A valid index of the collection. + /// + /// - Returns: The index immediately before `i`. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public func index(before index: Index) -> Index { + self._base.index(before: index) + } + + /// Offsets the given index by the specified distance. + /// + /// The value passed as `distance` must not offset `i` beyond the bounds of + /// the collection. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public func formIndex(_ i: inout Index, offsetBy distance: Int) { + self._base.formIndex(&i, offsetBy: distance) + } + + /// Returns an index that is the specified distance from the given index. + /// + /// The value passed as `distance` must not offset `i` beyond the bounds of + /// the collection, or the returned value will not be a valid index. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// + /// - Returns: An index offset by `distance` from the index `i`. If `distance` + /// is positive, this is the same value as the result of `distance` calls to + /// `index(after:)`. If `distance` is negative, this is the same value as + /// the result of `abs(distance)` calls to `index(before:)`. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public func index(_ i: Index, offsetBy distance: Int) -> Index { + self._base.index(i, offsetBy: distance) + } +} + +extension SortedDictionary.Keys { + /// Accesses the element at the specified position. + /// + /// - Parameter index: The position of the element to access. `index` must be + /// greater than or equal to `startIndex` and less than `endIndex`. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public subscript(position: Index) -> Key { + self._base[position].key + } +} + +extension SortedDictionary.Keys: Equatable { + /// Returns a Boolean value indicating whether two values are equal. + /// + /// Equality is the inverse of inequality. For any values `a` and `b`, + /// `a == b` implies that `a != b` is false. + /// + /// - Parameters: + /// - lhs: A value to compare. + /// - rhs: Another value to compare. + /// - Complexity: O(`self.count`) + @inlinable + public static func ==(lhs: Self, rhs: Self) -> Bool { + if lhs.count != rhs.count { return false } + for (e1, e2) in zip(lhs, rhs) { + if e1 == e2 { + return false + } + } + return true + } +} + +extension SortedDictionary.Keys: Hashable where Key: Hashable { + /// Hashes the essential components of this value by feeding them + /// into the given hasher. + /// - Parameter hasher: The hasher to use when combining + /// the components of this instance. + /// - Complexity: O(`self.count`) + @inlinable + public func hash(into hasher: inout Hasher) { + hasher.combine(self.count) + for key in self { + hasher.combine(key) + } + } +} diff --git a/Sources/SortedCollections/SortedDictionary/SortedDictionary+Partial RangeReplaceableCollection.swift b/Sources/SortedCollections/SortedDictionary/SortedDictionary+Partial RangeReplaceableCollection.swift new file mode 100644 index 000000000..e4a3e5f1b --- /dev/null +++ b/Sources/SortedCollections/SortedDictionary/SortedDictionary+Partial RangeReplaceableCollection.swift @@ -0,0 +1,174 @@ +//===----------------------------------------------------------------------===// +// +// 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 { + /// Returns a new sorted dictionary containing the key-value pairs of the + /// dictionary that satisfy the given predicate. + /// - Complexity: O(`n log n`) where `n` is the number of key-value pairs in the + /// sorted dictionary. + @inlinable + @inline(__always) + public func filter( + _ isIncluded: (Element) throws -> Bool + ) rethrows -> SortedDictionary { + let newTree: _Tree = try self._root.filter(isIncluded) + return SortedDictionary(_rootedAt: newTree) + } + + /// Removes and returns the first element of the collection. + /// + /// Calling this method may invalidate all saved indices of this collection. Do not rely on a + /// previously stored index value after altering a collection with any operation that can change + /// its length. + /// + /// - Returns: The first element of the collection if the collection is not empty; otherwise, nil. + /// - Complexity: O(`log n`) where `n` is the number of key-value pairs in the + /// sorted dictionary. + @inlinable + @inline(__always) + public mutating func popFirst() -> Element? { + self._root.popFirst() + } + + /// Removes and returns the last element of the collection. + /// + /// Calling this method may invalidate all saved indices of this collection. Do not rely on a + /// previously stored index value after altering a collection with any operation that can change + /// its length. + /// + /// - Returns: The last element of the collection if the collection is not empty; otherwise, nil. + /// - Complexity: O(`log n`) where `n` is the number of key-value pairs in the + /// sorted dictionary. + @inlinable + @inline(__always) + public mutating func popLast() -> Element? { + self._root.popLast() + } + + /// Removes and returns the first element of the collection. + /// + /// The collection must not be empty. + /// + /// Calling this method may invalidate all saved indices of this collection. Do not rely on a + /// previously stored index value after altering a collection with any operation that can change + /// its length. + /// + /// - Returns: The first element of the collection if the collection is not empty; otherwise, nil. + /// - Complexity: O(`log n`) where `n` is the number of key-value pairs in the + /// sorted dictionary. + @inlinable + @inline(__always) + @discardableResult + public mutating func removeFirst() -> Element { + self._root.removeFirst() + } + + /// Removes and returns the last element of the collection. + /// + /// The collection must not be empty. + /// + /// Calling this method may invalidate all saved indices of this collection. Do not rely on a + /// previously stored index value after altering a collection with any operation that can change + /// its length. + /// + /// - Returns: The last element of the collection if the collection is not empty; otherwise, nil. + /// - Complexity: O(`log n`) where `n` is the number of key-value pairs in the + /// sorted dictionary. + @inlinable + @inline(__always) + @discardableResult + public mutating func removeLast() -> Element { + self._root.removeLast() + } + + /// Removes the specified number of elements from the beginning of the collection. + /// + /// Calling this method may invalidate all saved indices of this collection. Do not rely on a + /// previously stored index value after altering a collection with any operation that can change + /// its length. + /// + /// - Parameter k: The number of elements to remove from the collection. `k` must be greater + /// than or equal to zero and must not exceed the number of elements in the collection. + /// - Complexity: O(`k log n`) where `n` is the number of key-value pairs in the + /// sorted dictionary. + @inlinable + @inline(__always) + public mutating func removeFirst(_ k: Int) { + self._root.removeFirst(k) + } + + /// Removes the specified number of elements from the end of the collection. + /// + /// Calling this method may invalidate all saved indices of this collection. Do not rely on a + /// previously stored index value after altering a collection with any operation that can change + /// its length. + /// + /// - Parameter k: The number of elements to remove from the collection. `k` must be greater + /// than or equal to zero and must not exceed the number of elements in the collection. + /// - Complexity: O(`k log n`) where `n` is the number of key-value pairs in the + /// sorted dictionary. + @inlinable + @inline(__always) + public mutating func removeLast(_ k: Int) { + self._root.removeLast(k) + } + + + /// Removes and returns the key-value pair at the specified index. + /// + /// Calling this method invalidates any existing indices for use with this sorted dictionary. + /// + /// - Parameter index: The position of the key-value pair to remove. `index` + /// must be a valid index of the sorted dictionary, and must not equal the sorted + /// dictionary’s end index. + /// - Returns: The key-value pair that correspond to `index`. + /// - Complexity: O(`log n`) where `n` is the number of key-value pairs in the + /// sorted dictionary. + @inlinable + @inline(__always) + public mutating func remove(at index: Index) -> Element { + index._index.ensureValid(forTree: self._root) + return self._root.remove(at: index._index) + } + + /// Removes the specified subrange of elements from the collection. + /// + /// - Parameter bounds: The subrange of the collection to remove. The bounds of the + /// range must be valid indices of the collection. + /// - Returns: The key-value pair that correspond to `index`. + /// - Complexity: O(`m log n`) where `n` is the number of key-value pairs in the + /// sorted dictionary, and `m` is the size of `bounds` + @inlinable + @inline(__always) + internal mutating func removeSubrange( + _ bounds: R + ) where R.Bound == Index { + let bounds = bounds.relative(to: self) + + bounds.upperBound._index.ensureValid(forTree: self._root) + bounds.lowerBound._index.ensureValid(forTree: self._root) + + return self._root.removeSubrange(Range(uncheckedBounds: (bounds.lowerBound._index, bounds.upperBound._index))) + } + + /// Removes all key-value pairs from the dictionary. + /// + /// Calling this method invalidates all indices with respect to the dictionary. + /// + /// - Complexity: O(`n`) + @inlinable + @inline(__always) + public mutating func removeAll() { + self._root.removeAll() + } + +} + diff --git a/Sources/SortedCollections/SortedDictionary/SortedDictionary+Sequence.swift b/Sources/SortedCollections/SortedDictionary/SortedDictionary+Sequence.swift new file mode 100644 index 000000000..ea4042edf --- /dev/null +++ b/Sources/SortedCollections/SortedDictionary/SortedDictionary+Sequence.swift @@ -0,0 +1,47 @@ +//===----------------------------------------------------------------------===// +// +// 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: Sequence { + @inlinable + public func forEach(_ body: (Element) throws -> Void) rethrows { + try self._root.forEach(body) + } + + /// An iterator over the elements of the sorted dictionary + public struct Iterator: IteratorProtocol { + @usableFromInline + internal var _iterator: _Tree.Iterator + + @inlinable + @inline(__always) + internal init(_base: SortedDictionary) { + self._iterator = _base._root.makeIterator() + } + + /// Advances to the next element and returns it, or nil if no next element exists. + /// + /// - Returns: The next element in the underlying sequence, if a next element exists; + /// otherwise, `nil`. + /// - Complexity: O(1) amortized over the entire sequence. + @inlinable + public mutating func next() -> Element? { + return self._iterator.next() + } + } + + /// Returns an iterator over the elements of the sorted dictionary. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + public __consuming func makeIterator() -> Iterator { + return Iterator(_base: self) + } +} diff --git a/Sources/SortedCollections/SortedDictionary/SortedDictionary+SubSequence.swift b/Sources/SortedCollections/SortedDictionary/SortedDictionary+SubSequence.swift new file mode 100644 index 000000000..9fbaee5a7 --- /dev/null +++ b/Sources/SortedCollections/SortedDictionary/SortedDictionary+SubSequence.swift @@ -0,0 +1,315 @@ +//===----------------------------------------------------------------------===// +// +// 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 { + public struct SubSequence { + @usableFromInline + internal typealias _TreeSubSequence = _Tree.SubSequence + + @usableFromInline + internal let _subSequence: _TreeSubSequence + + @inlinable + @inline(__always) + internal init(_ _subSequence: _TreeSubSequence) { + self._subSequence = _subSequence + } + + /// The underlying collection of the subsequence. + @inlinable + @inline(__always) + internal var base: SortedDictionary { SortedDictionary(_rootedAt: _subSequence.base) } + } +} + +extension SortedDictionary.SubSequence: Sequence { + public typealias Element = SortedDictionary.Element + + + public struct Iterator: IteratorProtocol { + @usableFromInline + internal var _iterator: _TreeSubSequence.Iterator + + @inlinable + @inline(__always) + internal init(_ _iterator: _TreeSubSequence.Iterator) { + self._iterator = _iterator + } + + /// Advances to the next element and returns it, or nil if no next element exists. + /// + /// - Returns: The next element in the underlying sequence, if a next element exists; + /// otherwise, `nil`. + /// - Complexity: O(1) amortized over the entire sequence. + @inlinable + @inline(__always) + public mutating func next() -> Element? { + _iterator.next() + } + } + + /// Returns an iterator over the elements of the subsequence. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public __consuming func makeIterator() -> Iterator { + Iterator(_subSequence.makeIterator()) + } +} + +extension SortedDictionary.SubSequence: BidirectionalCollection { + public typealias Index = SortedDictionary.Index + public typealias SubSequence = Self + + /// The position of the first element in a nonempty subsequence. + /// + /// If the collection is empty, `startIndex` is equal to `endIndex`. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var startIndex: Index { Index(_subSequence.startIndex) } + + /// The subsequence's "past the end" position---that is, the position one + /// greater than the last valid subscript argument. + /// + /// If the collection is empty, `endIndex` is equal to `startIndex`. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var endIndex: Index { Index(_subSequence.endIndex) } + + /// The number of elements in the subsequence. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var count: Int { _subSequence.count } + + /// Returns the distance between two indices. + /// + /// - 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. + /// - Complexity: O(1) + @inlinable + @inline(__always) + public func distance(from start: Index, to end: Index) -> Int { + start._index.ensureValid(forTree: _subSequence.base) + end._index.ensureValid(forTree: _subSequence.base) + return _subSequence.distance(from: start._index, to: end._index) + } + + + /// Returns the position immediately after the given index. + /// + /// - Parameter i: A valid index of the collection. `i` must be less than `endIndex`. + /// - Returns: The index value immediately after `i`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + @inline(__always) + public func index(after i: Index) -> Index { + i._index.ensureValid(forTree: _subSequence.base) + return Index(_subSequence.index(after: i._index)) + } + + /// Replaces the given index with its successor. + /// + /// - Parameter i: A valid index of the collection. `i` must be less than `endIndex`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + @inline(__always) + public func formIndex(after i: inout Index) { + i._index.ensureValid(forTree: _subSequence.base) + return _subSequence.formIndex(after: &i._index) + } + + + /// Returns the position immediately before the given index. + /// + /// - Parameter i: A valid index of the collection. `i` must be greater + /// than `startIndex`. + /// - Returns: The index value immediately before `i`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + @inline(__always) + public func index(before i: Index) -> Index { + i._index.ensureValid(forTree: _subSequence.base) + return Index(_subSequence.index(before: i._index)) + } + + /// Replaces the given index with its predecessor. + /// + /// - Parameter i: A valid index of the collection. `i` must be greater + /// than `startIndex`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + @inline(__always) + public func formIndex(before i: inout Index) { + i._index.ensureValid(forTree: _subSequence.base) + _subSequence.formIndex(before: &i._index) + } + + + /// Returns an index that is the specified distance from the given index. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - Returns: An index offset by `distance` from the index `i`. If `distance` + /// is positive, this is the same value as the result of `distance` calls to + /// `index(after:)`. If `distance` is negative, this is the same value as the + /// result of `abs(distance)` calls to `index(before:)`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + @inline(__always) + public func index(_ i: Index, offsetBy distance: Int) -> Index { + i._index.ensureValid(forTree: _subSequence.base) + return Index(_subSequence.index(i._index, offsetBy: distance)) + } + + /// Offsets the given index by the specified distance. + /// + /// The value passed as distance must not offset i beyond the bounds of the collection. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + @inline(__always) + internal func formIndex(_ i: inout Index, offsetBy distance: Int) { + i._index.ensureValid(forTree: _subSequence.base) + _subSequence.formIndex(&i._index, offsetBy: distance) + } + + + /// Returns an index that is the specified distance from the given index, unless that distance is beyond + /// a given limiting index. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - limit: A valid index of the collection to use as a limit. If `distance > 0`, a limit that is less + /// than `i` has no effect. Likewise, if `distance < 0`, a limit that is greater than `i` has + /// no effect. + /// - Returns: An index offset by `distance` from the index `i`, unless that index would be + /// beyond `limit` in the direction of movement. In that case, the method returns `nil`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + @inline(__always) + public func index(_ i: Index, offsetBy distance: Int, limitedBy limit: Index) -> Index? { + i._index.ensureValid(forTree: _subSequence.base) + limit._index.ensureValid(forTree: _subSequence.base) + + if let i = _subSequence.index(i._index, offsetBy: distance, limitedBy: limit._index) { + return Index(i) + } else { + return nil + } + } + + /// Offsets the given index by the specified distance, or so that it equals the given limiting index. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - limit: A valid index of the collection to use as a limit. If `distance > 0`, a limit that is less + /// than `i` has no effect. Likewise, if `distance < 0`, a limit that is greater than `i` has + /// no effect. + /// - Returns: `true` if `i` has been offset by exactly `distance` steps without going beyond + /// `limit`; otherwise, `false`. When the return value is `false`, the value of `i` is + /// equal to `limit`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + @inline(__always) + internal func formIndex(_ i: inout Index, offsetBy distance: Int, limitedBy limit: Self.Index) -> Bool { + i._index.ensureValid(forTree: _subSequence.base) + limit._index.ensureValid(forTree: _subSequence.base) + return _subSequence.formIndex(&i._index, offsetBy: distance, limitedBy: limit._index) + } + + @inlinable + @inline(__always) + public subscript(position: Index) -> Element { + position._index.ensureValid(forTree: _subSequence.base) + return _subSequence[position._index] + } + + @inlinable + public subscript(bounds: Range) -> SubSequence { + bounds.lowerBound._index.ensureValid(forTree: _subSequence.base) + bounds.upperBound._index.ensureValid(forTree: _subSequence.base) + + let bound = bounds.lowerBound._index..) { + _subSequence._failEarlyRangeCheck( + index._index, + bounds: bounds.lowerBound._index.., bounds: Range) { + _subSequence._failEarlyRangeCheck( + range.lowerBound._index.. Bool { + if lhs.count != rhs.count { return false } + for (e1, e2) in zip(lhs, rhs) { + if e1 == e2 { + return false + } + } + return true + } +} + +extension SortedDictionary.SubSequence: Hashable where Key: Hashable, Value: Hashable { + /// Hashes the essential components of this value by feeding them + /// into the given hasher. + /// - Parameter hasher: The hasher to use when combining + /// the components of this instance. + /// - Complexity: O(`self.count`) + @inlinable + public func hash(into hasher: inout Hasher) { + hasher.combine(self.count) + for (key, value) in self { + hasher.combine(key) + hasher.combine(value) + } + } +} diff --git a/Sources/SortedCollections/SortedDictionary/SortedDictionary+Subscripts.swift b/Sources/SortedCollections/SortedDictionary/SortedDictionary+Subscripts.swift new file mode 100644 index 000000000..e23de48b4 --- /dev/null +++ b/Sources/SortedCollections/SortedDictionary/SortedDictionary+Subscripts.swift @@ -0,0 +1,176 @@ +//===----------------------------------------------------------------------===// +// +// 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 { + /// Accesses the value associated with the key for both read and write operations + /// + /// This key-based subscript returns the value for the given key if the key is found in + /// the dictionary, or nil if the key is not found. + /// + /// When you assign a value for a key and that key already exists, the dictionary overwrites + /// the existing value. If the dictionary doesn’t contain the key, the key and value are added + /// as a new key-value pair. + /// + /// - Parameter key: The key to find in the dictionary. + /// - Returns: The value associated with key if key is in the dictionary; otherwise, nil. + /// - Complexity: O(`log n`) + @inlinable + @inline(__always) + public subscript(key: Key) -> Value? { + get { + return self._root.findAnyValue(forKey: key) + } + + _modify { + var (cursor, found) = self._root.takeCursor(forKey: key) + + var value: Value? + if found { + value = cursor.moveValue() + } + + defer { + if found { + if let value = value { + cursor.initializeValue(to: value) + } else { + cursor.removeElement(hasValueHole: true) + } + } else { + if let value = value { + cursor.insertElement( + (key, value), + capacity: self._root.internalCapacity + ) + } else { + // no-op + } + } + + cursor.apply(to: &self._root) + } + + yield &value + } + + set { + if let newValue = newValue { + self._root.updateAnyValue(newValue, forKey: key) + } else { + self._root.removeAnyElement(forKey: key) + } + } + } + + /// Accesses the value with the given key. If the dictionary doesn’t contain the given + /// key, accesses the provided default value as if the key and default value existed + /// in the dictionary. + /// + /// Use this subscript when you want either the value for a particular key or, when that + /// key is not present in the dictionary, a default value. + /// + /// - Parameters: + /// - key: The key the look up in the dictionary. + /// - defaultValue: The default value to use if key doesn’t exist in the dictionary. + /// - Returns: The value associated with key in the dictionary; otherwise, defaultValue. + /// - Complexity: O(`log n`) + @inlinable + @inline(__always) + public subscript( + key: Key, + default defaultValue: @autoclosure () -> Value + ) -> Value { + get { + return self[key] ?? defaultValue() + } + + set { + self[key] = newValue + } + + _modify { + var (cursor, found) = self._root.takeCursor(forKey: key) + + var value: Value + if found { + value = cursor.moveValue() + } else { + value = defaultValue() + } + + defer { + if found { + cursor.initializeValue(to: value) + } else { + cursor.insertElement( + (key, value), + capacity: self._root.internalCapacity + ) + } + + cursor.apply(to: &self._root) + } + + yield &value + } + + } + + /// Accesses the key-value pair at the specified position. + /// + /// This subscript takes an index into the sorted dictionary, instead of a key, and + /// returns the corresponding key-value pair as a tuple. When performing + /// collection-based operations that return an index into a dictionary, use this + /// subscript with the resulting value. + /// + /// For example, to find the key for a particular value in a sorted dictionary, use + /// the `firstIndex(where:)` method. + /// + /// let countryCodes: SortedDictionary = ["BR": "Brazil", "GH": "Ghana", "JP": "Japan"] + /// if let index = countryCodes.firstIndex(where: { $0.value == "Japan" }) { + /// print(countryCodes[index]) + /// print("Japan's country code is '\(countryCodes[index].key)'.") + /// } else { + /// print("Didn't find 'Japan' as a value in the dictionary.") + /// } + /// // Prints "(key: "JP", value: "Japan")" + /// // Prints "Japan's country code is 'JP'." + /// + /// - Parameter position: The position of the key-value pair to access. + /// `position` must be a valid index of the sorted dictionary and not equal + /// to `endIndex`. + /// - Returns: A two-element tuple with the key and value corresponding to + /// `position`. + /// - Complexity: O(1) + @inlinable + public subscript(position: Index) -> Element { + get { + position._index.ensureValid(forTree: self._root) + return self._root[position._index] + } + } + + /// Accesses a contiguous subrange of the collection's elements. + /// + /// - Parameter bounds: A range of the collection's indices. The bounds of + /// the range must be valid indices of the collection. + /// + /// - Complexity: O(1) + @inlinable + public subscript(bounds: Range) -> SubSequence { + bounds.lowerBound._index.ensureValid(forTree: _root) + bounds.upperBound._index.ensureValid(forTree: _root) + + let bound = bounds.lowerBound._index.. Element? { + _iterator.next()?.value + } + } + + @inlinable + @inline(__always) + public __consuming func makeIterator() -> Iterator { + Iterator(_base.makeIterator()) + } +} + +extension SortedDictionary.Values: BidirectionalCollection { + /// The index type for a dictionary's values view. + public typealias Index = SortedDictionary.Index + + /// The number of elements in the collection. + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var count: Int { self._base.count } + + /// A Boolean value that indicates whether the collection is empty. + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var isEmpty: Bool { self._base.isEmpty } + + /// The position of the first element in a nonempty collection. + /// + /// If the collection is empty, `startIndex` is equal to `endIndex`. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public var startIndex: Index { self._base.startIndex } + + /// The collection's "past the end" position---that is, the position one + /// greater than the last valid subscript argument. + /// + /// If the collection is empty, `endIndex` is equal to `startIndex`. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var endIndex: Index { self._base.endIndex } + + /// Returns the distance between two indices. + /// + /// - 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`. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public func distance(from start: Index, to end: Index) -> Int { + self._base.distance(from: start, to: end) + } + + /// Replaces the given index with its successor. + /// + /// The specified index must be a valid index less than `endIndex`, or the + /// returned value won't be a valid index in the collection. + /// + /// - Parameter i: A valid index of the collection. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public func formIndex(after index: inout Index) { + self._base.formIndex(after: &index) + } + + /// Returns the position immediately after the given index. + /// + /// The specified index must be a valid index less than `endIndex`, or the + /// returned value won't be a valid index in the collection. + /// + /// - Parameter i: A valid index of the collection. + /// + /// - Returns: The index immediately after `i`. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public func index(after index: Index) -> Index { + self._base.index(after: index) + } + + /// Replaces the given index with its successor. + /// + /// The specified index must be a valid index less than `endIndex`, or the + /// returned value won't be a valid index in the collection. + /// + /// - Parameter i: A valid index of the collection. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public func formIndex(before index: inout Index) { + self._base.formIndex(before: &index) + } + + /// Returns the position immediately before the given index. + /// + /// The specified index must be a valid index greater than `startIndex`, or + /// the returned value won't be a valid index in the collection. + /// + /// - Parameter i: A valid index of the collection. + /// + /// - Returns: The index immediately before `i`. + /// + /// - Complexity: O(`log n`) where `n` is the number of key-value pairs in the + /// sorted dictionary. + @inlinable + @inline(__always) + public func index(before index: Index) -> Index { + self._base.index(before: index) + } + + /// Offsets the given index by the specified distance. + /// + /// The value passed as `distance` must not offset `i` beyond the bounds of + /// the collection. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public func formIndex(_ i: inout Index, offsetBy distance: Int) { + self._base.formIndex(&i, offsetBy: distance) + } + + /// Returns an index that is the specified distance from the given index. + /// + /// The value passed as `distance` must not offset `i` beyond the bounds of + /// the collection, or the returned value will not be a valid index. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// + /// - Returns: An index offset by `distance` from the index `i`. If `distance` + /// is positive, this is the same value as the result of `distance` calls to + /// `index(after:)`. If `distance` is negative, this is the same value as + /// the result of `abs(distance)` calls to `index(before:)`. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public func index(_ i: Index, offsetBy distance: Int) -> Index { + self._base.index(i, offsetBy: distance) + } +} + +extension SortedDictionary.Values { + /// Accesses the element at the specified position. + /// + /// - Parameter index: The position of the element to access. `index` must be + /// greater than or equal to `startIndex` and less than `endIndex`. + /// + /// - Complexity: O(1) + @inlinable + public subscript(position: Index) -> Value { + get { + self._base[position].value + } + + _modify { + position._index.ensureValid(forTree: self._base._root) + + // Ensure we don't attempt to dereference the endIndex + precondition(position != endIndex, "Attempt to subscript out of range index.") + + var cursor = self._base._root.takeCursor(at: position._index) + var value = cursor.moveValue() + + defer { + cursor.initializeValue(to: value) + cursor.apply(to: &self._base._root) + } + + yield &value + } + } +} + +extension SortedDictionary.Values: Equatable where Value: Equatable { + /// Returns a Boolean value indicating whether two values are equal. + /// + /// Equality is the inverse of inequality. For any values `a` and `b`, + /// `a == b` implies that `a != b` is false. + /// + /// - Parameters: + /// - lhs: A value to compare. + /// - rhs: Another value to compare. + /// - Complexity: O(`self.count`) + @inlinable + public static func ==(lhs: Self, rhs: Self) -> Bool { + if lhs.count != rhs.count { return false } + for (e1, e2) in zip(lhs, rhs) { + if e1 == e2 { + return false + } + } + return true + } +} + +extension SortedDictionary.Values: Hashable where Value: Hashable { + /// Hashes the essential components of this value by feeding them + /// into the given hasher. + /// - Parameter hasher: The hasher to use when combining + /// the components of this instance. + /// - Complexity: O(`self.count`) + @inlinable + public func hash(into hasher: inout Hasher) { + hasher.combine(self.count) + for key in self { + hasher.combine(key) + } + } +} diff --git a/Sources/SortedCollections/SortedDictionary/SortedDictionary.Index.swift b/Sources/SortedCollections/SortedDictionary/SortedDictionary.Index.swift new file mode 100644 index 000000000..a86821101 --- /dev/null +++ b/Sources/SortedCollections/SortedDictionary/SortedDictionary.Index.swift @@ -0,0 +1,53 @@ +//===----------------------------------------------------------------------===// +// +// 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 { + /// Returns the index for a given key, if it exists + /// - Complexity: O(`log n`) + @inlinable + public func index(forKey key: Key) -> Index? { + if let index = self._root.findAnyIndex(forKey: key) { + return Index(index) + } else { + return nil + } + } + + /// The position of an element within a sorted dictionary + public struct Index { + @usableFromInline + internal var _index: _Tree.Index + + @inlinable + @inline(__always) + internal init(_ _index: _Tree.Index) { + self._index = _index + } + } +} + +// MARK: Equatable +extension SortedDictionary.Index: Equatable { + @inlinable + public static func ==(lhs: SortedDictionary.Index, rhs: SortedDictionary.Index) -> Bool { + lhs._index.ensureValid(with: rhs._index) + return lhs._index == rhs._index + } +} + +// MARK: Comparable +extension SortedDictionary.Index: Comparable { + @inlinable + public static func <(lhs: SortedDictionary.Index, rhs: SortedDictionary.Index) -> Bool { + lhs._index.ensureValid(with: rhs._index) + return lhs._index < rhs._index + } +} diff --git a/Sources/SortedCollections/SortedDictionary/SortedDictionary.swift b/Sources/SortedCollections/SortedDictionary/SortedDictionary.swift new file mode 100644 index 000000000..5c72e9d64 --- /dev/null +++ b/Sources/SortedCollections/SortedDictionary/SortedDictionary.swift @@ -0,0 +1,221 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +/// A collection which maintains key-value pairs in ascending sorted order. +/// +/// A sorted dictionary is a type of tree, providing efficient read and write operations +/// to the entries it contains. Each entry in the sorted dictionary is identified using a +/// key, which is a comparable type such as a string or number. You can use that key +/// to retrieve the corresponding value. +public struct SortedDictionary { + /// An element of the sorted dictionary. A key-value tuple. + public typealias Element = (key: Key, value: Value) + + @usableFromInline + internal typealias _Tree = _BTree + + @usableFromInline + internal var _root: _Tree + + /// Creates an empty dictionary. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public init() { + self._root = _Tree() + } + + /// Creates a dictionary rooted at a given BTree. + @inlinable + internal init(_rootedAt tree: _Tree) { + self._root = tree + } +} + +// MARK: Accessing Keys and Values +extension SortedDictionary { + /// A read-only collection view for the keys contained in this dictionary, as + /// an `SortedSet`. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var keys: Keys { Keys(_base: self) } + + /// A mutable collection view containing the values in this dictionary. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var values: Values { + get { + Values(_base: self) + } + + _modify { + let dummyDict = SortedDictionary(_rootedAt: _BTree.dummy) + var values = Values(_base: dummyDict) + swap(&self, &values._base) + defer { self = values._base } + yield &values + } + } +} + +// MARK: Mutations +extension SortedDictionary { + /// Ensures that the specified key exists in the dictionary (by appending one + /// with the supplied default value if necessary), then calls `body` to update + /// it in place. + /// + /// You can use this method to perform in-place operations on values in the + /// dictionary, whether or not `Value` has value semantics. The following + /// example uses this method while counting the occurrences of each letter + /// in a string: + /// + /// let message = "Hello, Elle!" + /// var letterCounts: SortedDictionary = [:] + /// for letter in message { + /// letterCounts.modifyValue(forKey: letter, default: 0) { count in + /// count += 1 + /// } + /// } + /// // letterCounts == ["H": 1, "e": 2, "l": 4, "o": 1, ...] + /// + /// - Parameters: + /// - key: The key to look up. If `key` does not already exist + /// in the dictionary, it is inserted with the supplied default value. + /// - defaultValue: The default value to append if `key` doesn't exist in + /// the dictionary. + /// - body: A function that performs an in-place mutation on the dictionary + /// value. + /// + /// - Returns: The return value of `body`. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + public mutating func modifyValue( + forKey key: Key, + default defaultValue: @autoclosure () -> Value, + _ body: (inout Value) throws -> R + ) rethrows -> R { + var (cursor, found) = self._root.takeCursor(forKey: key) + let r: R + + if found { + r = try cursor.updateCurrentNode { handle, slot in + try body(&handle[valueAt: slot]) + } + } else { + var value = defaultValue() + r = try body(&value) + cursor.insertElement((key, value), capacity: self._root.internalCapacity) + } + + cursor.apply(to: &self._root) + + return r + } + + + /// Updates the value stored in the dictionary for the given key, or + /// adds a new key-value pair if the key does not exist. + /// + /// Use this method instead of key-based subscripting when you need to + /// know whether the new value supplants the value of an existing key. If + /// the value of an existing key is updated, `updateValue(_:forKey:)` + /// returns the original value. + /// + /// - Parameters: + /// - value: The new value to add to the dictionary. + /// - key: The key to associate with value. If key already exists in the + /// dictionary, value replaces the existing associated value. If key + /// isn’t already a key of the dictionary, the `(key, value)` pair + /// is added. + /// - Returns: The value that was replaced, or nil if a new key-value + /// pair was added. + /// - Complexity: O(`log n`) where `n` is the number of key-value pairs in the + /// dictionary. + @inlinable + @discardableResult + public mutating func updateValue(_ value: Value, forKey key: Key) -> Value? { + self._root.updateAnyValue(value, forKey: key)?.value + } +} + +// MARK: Removing Keys and Values +extension SortedDictionary { + /// Removes the given key and its associated value from the sorted dictionary. + /// + /// Calling this method invalidates any existing indices for use with this sorted dictionary. + /// + /// - Parameter key: The key to remove along with its associated value. + /// - Returns: The value that was removed, or `nil` if the key was not present + /// in the dictionary. + /// - Complexity: O(`log n`) where `n` is the number of key-value pairs in the + /// dictionary. + @inlinable + @inline(__always) + public mutating func removeValue(forKey key: Key) -> Value? { + return self._root.removeAnyElement(forKey: key)?.value + } +} + +// MARK: Transforming a Dictionary +extension SortedDictionary { + /// Returns a new dictionary containing the keys of this dictionary with the values + /// transformed by the given closure. + /// + /// - Parameter transform: A closure that transforms a value. transform accepts + /// each value of the dictionary in order as its parameter and returns a transformed + /// value of the same or of a different type. + /// - Returns: A sorted dictionary containing the keys and transformed values of + /// this dictionary. + /// - Complexity: O(`n`) + @inlinable + @inline(__always) + public func mapValues( + _ transform: (Value) throws -> T + ) rethrows -> SortedDictionary { + let tree = try self._root.mapValues(transform) + return SortedDictionary(_rootedAt: tree) + } + + /// Returns a new dictionary containing only the key-value pairs that have non-nil values + /// as the result of transformation by the given closure. + /// + /// Use this method to receive a dictionary with non-optional values when your transformation + /// produces optional values. + /// + /// - Parameter transform: A closure that transforms a value. `transform` accepts + /// each value of the dictionary as its parameter and returns an optional transformed + /// value of the same or of a different type. + /// - Returns: A sorted dictionary containing the keys and non-nil transformed values + /// of this dictionary. + /// - Complexity: O(`n`) + func compactMapValues( + _ transform: (Value) throws -> T? + ) rethrows -> SortedDictionary { + // TODO: Optimize to reuse dictionary structure. + // TODO: optimize to use identify fastest iteration method + // TODO: optimize CoW checks + var builder = _BTree.Builder() + + for (key, value) in self { + if let newValue = try transform(value) { + builder.append((key, newValue)) + } + } + + return SortedDictionary(_rootedAt: builder.finish()) + } +} diff --git a/Sources/SortedCollections/SortedSet/SortedSet+BidirectionalCollection.swift b/Sources/SortedCollections/SortedSet/SortedSet+BidirectionalCollection.swift new file mode 100644 index 000000000..a108ff244 --- /dev/null +++ b/Sources/SortedCollections/SortedSet/SortedSet+BidirectionalCollection.swift @@ -0,0 +1,175 @@ +//===----------------------------------------------------------------------===// +// +// 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 SortedSet: BidirectionalCollection { + /// The number of elements in the sorted set. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var count: Int { self._root.count } + + /// A Boolean value that indicates whether the set is empty. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var isEmpty: Bool { self._root.isEmpty } + + /// The position of the first element in a nonempty set. + /// + /// If the collection is empty, `startIndex` is equal to `endIndex`. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public var startIndex: Index { Index(self._root.startIndex) } + + /// The set's "past the end" position---that is, the position one + /// greater than the last valid subscript argument. + /// + /// If the collection is empty, `endIndex` is equal to `startIndex`. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var endIndex: Index { Index(self._root.endIndex) } + + /// Returns the distance between two indices. + /// + /// - 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. + /// - Complexity: O(1) + @inlinable + public func distance(from start: Index, to end: Index) -> Int { + start._index.ensureValid(forTree: self._root) + end._index.ensureValid(forTree: self._root) + return self._root.distance(from: start._index, to: end._index) + } + + /// Replaces the given index with its successor. + /// + /// - Parameter index: A valid index of the collection. `index` must be less than `endIndex`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + public func formIndex(after index: inout Index) { + index._index.ensureValid(forTree: self._root) + self._root.formIndex(after: &index._index) + } + + /// Returns the position immediately after the given index. + /// + /// - Parameter index: A valid index of the collection. `index` must be less than `endIndex`. + /// - Returns: The index value immediately after `index`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + public func index(after index: Index) -> Index { + index._index.ensureValid(forTree: self._root) + return Index(self._root.index(after: index._index)) + } + + /// Replaces the given index with its predecessor. + /// + /// - Parameter index: A valid index of the collection. `index` must be greater + /// than `startIndex`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + public func formIndex(before index: inout Index) { + index._index.ensureValid(forTree: self._root) + self._root.formIndex(before: &index._index) + } + + /// Returns the position immediately before the given index. + /// + /// - Parameter index: A valid index of the collection. `index` must be greater + /// than `startIndex`. + /// - Returns: The index value immediately before `index`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + public func index(before index: Index) -> Index { + index._index.ensureValid(forTree: self._root) + return Index(self._root.index(before: index._index)) + } + + /// Offsets the given index by the specified distance. + /// + /// The value passed as distance must not offset i beyond the bounds of the collection. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + public func formIndex(_ i: inout Index, offsetBy distance: Int) { + i._index.ensureValid(forTree: self._root) + self._root.formIndex(&i._index, offsetBy: distance) + } + + /// Returns an index that is the specified distance from the given index. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - Returns: An index offset by `distance` from the index `i`. If `distance` + /// is positive, this is the same value as the result of `distance` calls to + /// `index(after:)`. If `distance` is negative, this is the same value as the + /// result of `abs(distance)` calls to `index(before:)`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + public func index(_ i: Index, offsetBy distance: Int) -> Index { + i._index.ensureValid(forTree: self._root) + return Index(self._root.index(i._index, offsetBy: distance)) + } + + /// Returns an index that is the specified distance from the given index, unless that distance is beyond + /// a given limiting index. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - limit: A valid index of the collection to use as a limit. If `distance > 0`, a limit that is less + /// than `i` has no effect. Likewise, if `distance < 0`, a limit that is greater than `i` has + /// no effect. + /// - Returns: An index offset by `distance` from the index `i`, unless that index would be + /// beyond `limit` in the direction of movement. In that case, the method returns `nil`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + public func index(_ i: Index, offsetBy distance: Int, limitedBy limit: Index) -> Index? { + i._index.ensureValid(forTree: self._root) + limit._index.ensureValid(forTree: self._root) + if let i = self._root.index(i._index, offsetBy: distance, limitedBy: limit._index) { + return Index(i) + } else { + return nil + } + } + + /// Offsets the given index by the specified distance, or so that it equals the given limiting index. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - limit: A valid index of the collection to use as a limit. If `distance > 0`, a limit that is less + /// than `i` has no effect. Likewise, if `distance < 0`, a limit that is greater than `i` has + /// no effect. + /// - Returns: `true` if `i` has been offset by exactly `distance` steps without going beyond + /// `limit`; otherwise, `false`. When the return value is `false`, the value of `i` is + /// equal to `limit`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + internal func formIndex(_ i: inout Index, offsetBy distance: Int, limitedBy limit: Self.Index) -> Bool { + i._index.ensureValid(forTree: self._root) + limit._index.ensureValid(forTree: self._root) + return self._root.formIndex(&i._index, offsetBy: distance, limitedBy: limit._index) + } +} diff --git a/Sources/SortedCollections/SortedSet/SortedSet+Codable.swift b/Sources/SortedCollections/SortedSet/SortedSet+Codable.swift new file mode 100644 index 000000000..6534e7f26 --- /dev/null +++ b/Sources/SortedCollections/SortedSet/SortedSet+Codable.swift @@ -0,0 +1,52 @@ +//===----------------------------------------------------------------------===// +// +// 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 SortedSet: Encodable where Element: Encodable { + /// Encodes the elements of this ordered set into the given encoder. + /// + /// - Parameter encoder: The encoder to write data to. + @inlinable + public func encode(to encoder: Encoder) throws { + var container = encoder.unkeyedContainer() + try self.forEach { element in + try container.encode(element) + } + } +} + +extension SortedSet: Decodable where Element: Decodable { + /// Creates a new ordered set by decoding from the given decoder. + /// + /// This initializer throws an error if reading from the decoder fails, or + /// if the decoded contents contain duplicate values. + /// + /// - Parameter decoder: The decoder to read data from. + @inlinable + public init(from decoder: Decoder) throws { + var container = try decoder.unkeyedContainer() + var builder = _Tree.Builder(deduplicating: true) + var previousElement: Element? = nil + + while !container.isAtEnd { + let element = try container.decode(Element.self) + guard previousElement == nil || previousElement! < element else { + let context = DecodingError.Context( + codingPath: container.codingPath, + debugDescription: "Decoded elements out of order.") + throw DecodingError.dataCorrupted(context) + } + builder.append(element) + previousElement = element + } + + self.init(_rootedAt: builder.finish()) + } +} diff --git a/Sources/SortedCollections/SortedSet/SortedSet+CustomReflectable.swift b/Sources/SortedCollections/SortedSet/SortedSet+CustomReflectable.swift new file mode 100644 index 000000000..4177fc614 --- /dev/null +++ b/Sources/SortedCollections/SortedSet/SortedSet+CustomReflectable.swift @@ -0,0 +1,17 @@ +//===----------------------------------------------------------------------===// +// +// 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 SortedSet: CustomReflectable { + /// The custom mirror for this instance. + public var customMirror: Mirror { + Mirror(self, unlabeledChildren: self, displayStyle: .set) + } +} diff --git a/Sources/SortedCollections/SortedSet/SortedSet+CustomStringConvertible.swift b/Sources/SortedCollections/SortedSet/SortedSet+CustomStringConvertible.swift new file mode 100644 index 000000000..eeade805f --- /dev/null +++ b/Sources/SortedCollections/SortedSet/SortedSet+CustomStringConvertible.swift @@ -0,0 +1,45 @@ +//===----------------------------------------------------------------------===// +// +// 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 SortedSet: CustomStringConvertible, CustomDebugStringConvertible { + @inlinable + public var description: String { + var result = "[" + var first = true + for element in self { + if first { + first = false + } else { + result += ", " + } + print(element, terminator: "", to: &result) + } + result += "]" + return result + } + + @inlinable + public var debugDescription: String { + var result = "SortedSet<\(Element.self)>([" + var first = true + for element in self { + if first { + first = false + } else { + result += ", " + } + + debugPrint(element, terminator: "", to: &result) + } + result += "])" + return result + } +} diff --git a/Sources/SortedCollections/SortedSet/SortedSet+Equatable.swift b/Sources/SortedCollections/SortedSet/SortedSet+Equatable.swift new file mode 100644 index 000000000..1818c1e10 --- /dev/null +++ b/Sources/SortedCollections/SortedSet/SortedSet+Equatable.swift @@ -0,0 +1,33 @@ +//===----------------------------------------------------------------------===// +// +// 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 SortedSet: Equatable { + /// Returns a Boolean value indicating whether two values are equal. + /// + /// Equality is the inverse of inequality. For any values `a` and `b`, + /// `a == b` implies that `a != b` is false. + /// + /// - Parameters: + /// - lhs: A value to compare. + /// - rhs: Another value to compare. + /// - Complexity: O(`self.count`) + @inlinable + public static func ==(lhs: Self, rhs: Self) -> Bool { + // TODO: optimize/benchmarking by comparing node identity. + if lhs.count != rhs.count { return false } + for (k1, k2) in zip(lhs, rhs) { + if k1 != k2 { + return false + } + } + return true + } +} diff --git a/Sources/SortedCollections/SortedSet/SortedSet+ExpressibleByArrayLiteral.swift b/Sources/SortedCollections/SortedSet/SortedSet+ExpressibleByArrayLiteral.swift new file mode 100644 index 000000000..01f2b040d --- /dev/null +++ b/Sources/SortedCollections/SortedSet/SortedSet+ExpressibleByArrayLiteral.swift @@ -0,0 +1,32 @@ +//===----------------------------------------------------------------------===// +// +// 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 SortedSet: ExpressibleByArrayLiteral { + /// Creates a new sorted set from the contents of an array literal. + /// + /// Duplicate elements in the literal are allowed, but the resulting + /// set will only contain the last occurrence of each. + /// + /// Do not call this initializer directly. It is used by the compiler when + /// you use an array literal. Instead, create a new set using an array + /// literal as its value by enclosing a comma-separated list of values in + /// square brackets. You can use an array literal anywhere a set is expected + /// by the type context. + /// + /// - Parameter elements: A variadic list of elements of the new set. + /// + /// - Complexity: O(`n log n`) where `n` is the number of elements + /// in `elements`. + @inlinable + public init(arrayLiteral elements: Element...) { + self.init(elements) + } +} diff --git a/Sources/SortedCollections/SortedSet/SortedSet+Hashable.swift b/Sources/SortedCollections/SortedSet/SortedSet+Hashable.swift new file mode 100644 index 000000000..d5cea514f --- /dev/null +++ b/Sources/SortedCollections/SortedSet/SortedSet+Hashable.swift @@ -0,0 +1,25 @@ +//===----------------------------------------------------------------------===// +// +// 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 SortedSet: Hashable where Element: Hashable { + /// Hashes the essential components of this value by feeding them + /// into the given hasher. + /// - Parameter hasher: The hasher to use when combining + /// the components of this instance. + /// - Complexity: O(`self.count`) + @inlinable + public func hash(into hasher: inout Hasher) { + hasher.combine(self.count) + for element in self { + hasher.combine(element) + } + } +} diff --git a/Sources/SortedCollections/SortedSet/SortedSet+Initializers.swift b/Sources/SortedCollections/SortedSet/SortedSet+Initializers.swift new file mode 100644 index 000000000..dae3a783e --- /dev/null +++ b/Sources/SortedCollections/SortedSet/SortedSet+Initializers.swift @@ -0,0 +1,53 @@ +//===----------------------------------------------------------------------===// +// +// 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 SortedSet { + /// Creates a new set from a finite sequence of items. + /// + /// - Parameter elements: The elements to use as members of the new set. + /// - Complexity: O(`n log n`) where `n` is the number of elements + /// in the sequence. + @inlinable + public init(_ elements: S) where S.Element == Element { + self.init() + + for element in elements { + self._root.updateAnyValue((), forKey: element) + } + } + + /// Creates a dictionary from a sequence of **sorted** elements. + /// + /// This is a more efficient alternative to ``init(_:)`` which offers + /// better asymptotic performance, and also reduces memory usage when constructing a + /// sorted set on a pre-sorted sequence. + /// + /// - Parameter elements: A sequence of elements in non-decreasing comparison order for the + /// new set. + /// - Complexity: O(`n`) where `n` is the number of elements in the + /// sequence. + @inlinable + public init( + sortedElements elements: S + ) where S.Element == Element { + var builder = _Tree.Builder() + + var previousElement: Element? = nil + for element in elements { + precondition(previousElement == nil || previousElement! < element, + "Sequence out of order.") + builder.append(element) + previousElement = element + } + + self.init(_rootedAt: builder.finish()) + } +} diff --git a/Sources/SortedCollections/SortedSet/SortedSet+Partial RangeReplaceableCollection.swift b/Sources/SortedCollections/SortedSet/SortedSet+Partial RangeReplaceableCollection.swift new file mode 100644 index 000000000..569f2faa5 --- /dev/null +++ b/Sources/SortedCollections/SortedSet/SortedSet+Partial RangeReplaceableCollection.swift @@ -0,0 +1,174 @@ +//===----------------------------------------------------------------------===// +// +// 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 SortedSet { + /// Returns a new sorted set containing the members of the set that satisfy the given + /// predicate. + /// - Complexity: O(`n log n`) where `n` is the number of members in the + /// sorted set. + @inlinable + @inline(__always) + public func filter( + _ isIncluded: (Element) throws -> Bool + ) rethrows -> SortedSet { + let newTree: _Tree = try self._root.filter({ try isIncluded($0.key) }) + return SortedSet(_rootedAt: newTree) + } + + /// Removes and returns the first element of the collection. + /// + /// Calling this method may invalidate all saved indices of this collection. Do not rely on a + /// previously stored index value after altering a collection with any operation that can change + /// its length. + /// + /// - Returns: The first element of the collection if the collection is not empty; otherwise, nil. + /// - Complexity: O(`log n`) where `n` is the number of members in the + /// sorted set. + @inlinable + @inline(__always) + public mutating func popFirst() -> Element? { + self._root.popFirst()?.key + } + + /// Removes and returns the last element of the collection. + /// + /// Calling this method may invalidate all saved indices of this collection. Do not rely on a + /// previously stored index value after altering a collection with any operation that can change + /// its length. + /// + /// - Returns: The last element of the collection if the collection is not empty; otherwise, nil. + /// - Complexity: O(`log n`) where `n` is the number of members in the + /// sorted set. + @inlinable + @inline(__always) + public mutating func popLast() -> Element? { + self._root.popLast()?.key + } + + /// Removes and returns the first element of the collection. + /// + /// The collection must not be empty. + /// + /// Calling this method may invalidate all saved indices of this collection. Do not rely on a + /// previously stored index value after altering a collection with any operation that can change + /// its length. + /// + /// - Returns: The first element of the collection if the collection is not empty; otherwise, nil. + /// - Complexity: O(`log n`) where `n` is the number of members in the + /// sorted set. + @inlinable + @inline(__always) + @discardableResult + public mutating func removeFirst() -> Element { + self._root.removeFirst().key + } + + /// Removes and returns the last element of the collection. + /// + /// The collection must not be empty. + /// + /// Calling this method may invalidate all saved indices of this collection. Do not rely on a + /// previously stored index value after altering a collection with any operation that can change + /// its length. + /// + /// - Returns: The last element of the collection if the collection is not empty; otherwise, nil. + /// - Complexity: O(`log n`) where `n` is the number of members in the + /// sorted set. + @inlinable + @inline(__always) + @discardableResult + public mutating func removeLast() -> Element { + self._root.removeLast().key + } + + /// Removes the specified number of elements from the beginning of the collection. + /// + /// Calling this method may invalidate all saved indices of this collection. Do not rely on a + /// previously stored index value after altering a collection with any operation that can change + /// its length. + /// + /// - Parameter k: The number of elements to remove from the collection. `k` must be greater + /// than or equal to zero and must not exceed the number of elements in the collection. + /// - Complexity: O(`k log n`) where `n` is the number of members in the + /// sorted set. + @inlinable + @inline(__always) + public mutating func removeFirst(_ k: Int) { + self._root.removeFirst(k) + } + + /// Removes the specified number of elements from the end of the collection. + /// + /// Calling this method may invalidate all saved indices of this collection. Do not rely on a + /// previously stored index value after altering a collection with any operation that can change + /// its length. + /// + /// - Parameter k: The number of elements to remove from the collection. `k` must be greater + /// than or equal to zero and must not exceed the number of elements in the collection. + /// - Complexity: O(`k log n`) where `n` is the number of members in the + /// sorted set. + @inlinable + @inline(__always) + public mutating func removeLast(_ k: Int) { + self._root.removeLast(k) + } + + + /// Removes and returns the key-value pair at the specified index. + /// + /// Calling this method invalidates any existing indices for use with this sorted dictionary. + /// + /// - Parameter index: The position of the key-value pair to remove. `index` + /// must be a valid index of the sorted dictionary, and must not equal the sorted + /// dictionary’s end index. + /// - Returns: The key-value pair that correspond to `index`. + /// - Complexity: O(`log n`) where `n` is the number of members in the + /// sorted set. + @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 + } + + /// Removes the specified subrange of elements from the collection. + /// + /// - Parameter bounds: The subrange of the collection to remove. The bounds of the + /// range must be valid indices of the collection. + /// - Returns: The key-value pair that correspond to `index`. + /// - Complexity: O(`m log n`) where `n` is the number of elements in the + /// sorted set, and `m` is the size of `bounds` + @inlinable + @inline(__always) + internal mutating func removeSubrange( + _ bounds: R + ) where R.Bound == Index { + // TODO: optimize to not perform excessive traversals + let bounds = bounds.relative(to: self) + + bounds.upperBound._index.ensureValid(forTree: self._root) + bounds.lowerBound._index.ensureValid(forTree: self._root) + + return self._root.removeSubrange(Range(uncheckedBounds: (bounds.lowerBound._index, bounds.upperBound._index))) + } + + /// Removes all elements from the set. + /// + /// Calling this method invalidates all indices with respect to the set. + /// + /// - Complexity: O(`n`) + @inlinable + @inline(__always) + public mutating func removeAll() { + self._root.removeAll() + } +} + diff --git a/Sources/SortedCollections/SortedSet/SortedSet+Sequence.swift b/Sources/SortedCollections/SortedSet/SortedSet+Sequence.swift new file mode 100644 index 000000000..e9f3bac9e --- /dev/null +++ b/Sources/SortedCollections/SortedSet/SortedSet+Sequence.swift @@ -0,0 +1,51 @@ +//===----------------------------------------------------------------------===// +// +// 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 SortedSet: Sequence { + @inlinable + @inline(__always) + public func forEach(_ body: (Element) throws -> Void) rethrows { + try self._root.forEach({ try body($0.key) }) + } + + /// An iterator over the elements of the sorted set + @frozen + public struct Iterator: IteratorProtocol { + @usableFromInline + internal var _iterator: _Tree.Iterator + + @inlinable + @inline(__always) + internal init(_base: SortedSet) { + self._iterator = _base._root.makeIterator() + } + + /// Advances to the next element and returns it, or nil if no next element exists. + /// + /// - Returns: The next element in the underlying sequence, if a next element exists; + /// otherwise, `nil`. + /// - Complexity: O(1) amortized over the entire sequence. + @inlinable + @inline(__always) + public mutating func next() -> Element? { + return self._iterator.next()?.key + } + } + + /// Returns an iterator over the elements of the sorted set. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public __consuming func makeIterator() -> Iterator { + return Iterator(_base: self) + } +} diff --git a/Sources/SortedCollections/SortedSet/SortedSet+SetAlgebra.swift b/Sources/SortedCollections/SortedSet/SortedSet+SetAlgebra.swift new file mode 100644 index 000000000..f05df3123 --- /dev/null +++ b/Sources/SortedCollections/SortedSet/SortedSet+SetAlgebra.swift @@ -0,0 +1,370 @@ +//===----------------------------------------------------------------------===// +// +// 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 SortedSet: SetAlgebra { + + // MARK: Testing for Membership + + /// Returns a Boolean value that indicates whether the given element exists in the set. + /// - Complexity: O(`log n`) where `n` is the number of members in the + /// sorted set. + @inlinable + @inline(__always) + public func contains(_ member: Element) -> Bool { + self._root.contains(key: member) + } + + + // MARK: Adding Elements + + /// Inserts the given element in the set if it is not already present. + /// + /// - Parameter newMember: + /// - Returns: `(true, newMember)` if `newMember` was not contained in the + /// set. If an element equal to `newMember` was already contained in the set, the + /// method returns `(false, oldMember)`, where `oldMember` is the element + /// that was equal to `newMember`. In some cases, `oldMember` may be + /// distinguishable from `newMember` by identity comparison or some other means. + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + @discardableResult + public mutating func insert( + _ newMember: Element + ) -> (inserted: Bool, memberAfterInsert: Element) { + if let oldKey = self._root.updateAnyValue((), forKey: newMember, updatingKey: false)?.key { + return (inserted: false, memberAfterInsert: oldKey) + } else { + return (inserted: true, memberAfterInsert: newMember) + } + } + + /// Inserts the given element into the set unconditionally. + /// + /// - Parameter newMember: An element to insert into the set. + /// - Returns: An element equal to `newMember` if the set already contained such a + /// member; otherwise, `nil`. In some cases, the returned element may be distinguishable + /// from `newMember` by identity comparison or some other means. + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + @discardableResult + public mutating func update(with newMember: Element) -> Element? { + return self._root.updateAnyValue((), forKey: newMember, updatingKey: true)?.key + } + + /// Removes the given element from the set. + /// + /// - Parameter member: The element of the set to remove. + /// + /// - Returns: The element equal to `member` if `member` is contained in the + /// set; otherwise, `nil`. In some cases, the returned element may be + /// distinguishable from `newMember` by identity comparison or some other + /// means. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + @discardableResult + public mutating func remove(_ member: Element) -> Element? { + return self._root.removeAnyElement(forKey: member)?.key + } + + // MARK: Combining Sets + /// Returns a new set with the elements of both this and the given set. + /// + /// - Parameter other: A sorted set of the same type as the current set. + /// - Returns: A new sorted set with the unique elements of this set and `other`. + /// - Note: if this set and `other` contain elements that are equal but + /// distinguishable (e.g. via `===`), the element from the second set is inserted. + /// - Complexity: O(`self.count` + `other.count`) + @inlinable + public func union(_ other: __owned Self) -> Self { + var builder = _Tree.Builder(deduplicating: true) + + var it1 = self.makeIterator() + var it2 = other.makeIterator() + + var e1 = it1.next() + var e2 = it2.next() + + // While both sequences have a value, consume the smallest element. + while let el1 = e1, let el2 = e2 { + if el1 < el2 { + builder.append(el1) + e1 = it1.next() + } else { + builder.append(el2) + e2 = it2.next() + } + } + + while let el1 = e1 { + builder.append(el1) + e1 = it1.next() + } + + while let el2 = e2 { + builder.append(el2) + e2 = it2.next() + } + + return SortedSet(_rootedAt: builder.finish()) + } + + /// Adds the elements of the given set to the set. + /// + /// - Parameter other: A set of the same type as the current set. + /// - Note: if this set and `other` contain elements that are equal but + /// distinguishable (e.g. via `===`), the element from the second set is inserted. + /// - Complexity: O(`self.count` + `other.count`) + @inlinable + public mutating func formUnion(_ other: __owned Self) { + self = union(other) + } + + /// Returns a new set with the elements that are common to both this set and + /// the given set. + /// + /// - Parameter other: A set of the same type as the current set. + /// - Returns: A new set. + /// - Note: if this set and `other` contain elements that are equal but + /// distinguishable (e.g. via `===`), which of these elements is present + /// in the result is unspecified. + /// - Complexity: O(`self.count` + `other.count`) + @inlinable + public func intersection(_ other: Self) -> Self { + // We'll run this such that 'self' is the smallest array + // TODO: might want to consider uniqueness to minimize CoW copies. + var builder = _Tree.Builder(deduplicating: true) + + var it1 = self.makeIterator() + var it2 = other.makeIterator() + + var e1 = it1.next() + var e2 = it2.next() + + // While both sequences have a value, consume the smallest element. + while let el1 = e1, let el2 = e2 { + if el1 < el2 { + e1 = it1.next() + } else if el1 == el2 { + builder.append(el2) + e1 = it1.next() + e2 = it2.next() + } else { + // el1 > el1 + e2 = it2.next() + } + } + + return SortedSet(_rootedAt: builder.finish()) + } + + /// Removes the elements of this set that aren't also in the given set. + /// + /// - Parameter other: A set of the same type as the current set. + /// - Note: if this set and `other` contain elements that are equal but + /// distinguishable (e.g. via `===`), which of these elements is present + /// in the result is unspecified. + /// - Complexity: O(`self.count` + `other.count`) + @inlinable + public mutating func formIntersection(_ other: Self) { + self = intersection(other) + } + + /// Returns a new set with the elements that are either in this set or in the + /// given set, but not in both. + /// + /// - Parameter other: A set of the same type as the current set. + /// - Returns: A new set. + /// - Complexity: O(`self.count` + `other.count`) + @inlinable + public func symmetricDifference(_ other: Self) -> Self { + var builder = _Tree.Builder(deduplicating: true) + + var it1 = self.makeIterator() + var it2 = other.makeIterator() + + var e1 = it1.next() + var e2 = it2.next() + + // While both sequences have a value, consume the smallest element. + while let el1 = e1, let el2 = e2 { + if el1 < el2 { + builder.append(el1) + e1 = it1.next() + } else if el2 < el1 { + builder.append(el2) + e2 = it2.next() + } else { + // e1 == e2 + e1 = it1.next() + e2 = it2.next() + } + } + + while let el1 = e1 { + builder.append(el1) + e1 = it1.next() + } + + while let el2 = e2 { + builder.append(el2) + e2 = it2.next() + } + + return SortedSet(_rootedAt: builder.finish()) + } + + /// Removes the elements of the set that are also in the given set and adds + /// the members of the given set that are not already in the set. + /// + /// - Parameter other: A set of the same type. + /// - Complexity: O(`self.count` + `other.count`) + @inlinable + public mutating func formSymmetricDifference(_ other: Self) { + self = self.symmetricDifference(other) + } + + /// Returns a new set containing the elements of this set that do not occur + /// in the given set. + /// + /// - Parameter other: A set of the same type as the current set. + /// - Returns: A new set. + /// - Complexity: O(`self.count` + `other.count`) + @inlinable + public func subtracting(_ other: Self) -> Self { + var builder = _Tree.Builder(deduplicating: true) + + var it1 = self.makeIterator() + var it2 = other.makeIterator() + + var e1 = it1.next() + var e2 = it2.next() + + // While both sequences have a value, consume the smallest element. + while let el1 = e1, let el2 = e2 { + if el1 < el2 { + builder.append(el1) + e1 = it1.next() + } else if el2 < el1 { + e2 = it2.next() + } else { + // e1 == e2 + e1 = it1.next() + e2 = it2.next() + } + } + + while let el1 = e1 { + builder.append(el1) + e1 = it1.next() + } + + return SortedSet(_rootedAt: builder.finish()) + } + + /// Removes the elements of the given set from this set. + /// + /// - Parameter other: A set of the same type as the current set. + /// - Complexity: O(`self.count` + `other.count`) + @inlinable + public mutating func subtract(_ other: SortedSet) { + self = self.subtracting(other) + } + + // MARK: Comparing Sets + /// Returns a Boolean value that indicates whether the set is a subset of another set. + /// + /// Set _A_ is a subset of another set _B_ if every member of _A_ is also a member of _B_. + /// + /// - Parameter other: A set of the same type as the current set. + /// - Returns: `true` if the set is a subset of other; otherwise, `false`. + /// - Complexity: O(max(`self.count`, `other.count`)) + @inlinable + public func isSubset(of other: SortedSet) -> Bool { + // TODO: could be worthwhile to evaluate recursive approach + // TODO: in some cases, it could be faster to search from the root each time + // Searching from the root is faster when: + // self.count < other.count / (log(other.count) - 1) + // This means when `other` is significantly larger than `self`, it may be + // faster to search from the root each time. + if self.count > other.count { return false } + + var superIterator = other.makeIterator() + + for element in self { + while true { + // If we exhausted the superset without finding our element, then it + // does not exist in the superset. + guard let superElement = superIterator.next() else { return false } + + // If the superElement is greater than element, we won't find element + // further on in the superset and therefore it doesn't exist in it. Here + // we can return false + if superElement > element { return false } + + // We did find the element + if superElement == element { break } + } + } + + return true + } + + /// Returns a Boolean value that indicates whether this set is a strict + /// subset of the given set. + /// + /// Set *A* is a strict subset of another set *B* if every member of *A* is + /// also a member of *B* and *B* contains at least one element that is not a + /// member of *A*. + /// + /// - Parameter other: A set of the same type as the current set. + /// - Returns: `true` if the set is a strict subset of `other`; otherwise, + /// `false`. + /// - Complexity: O(`self.count` + `other.count`). + @inlinable + public func isStrictSubset(of other: SortedSet) -> Bool { + if self.count >= other.count { return false } + return self.isSubset(of: other) + } + + /// Returns a Boolean value that indicates whether the set has no members in + /// common with the given set. + /// + /// - Parameter other: A set of the same type as the current set. + /// - Returns: `true` if the set has no elements in common with `other`; + /// otherwise, `false`. + /// - Complexity: O(`self.count` + `other.count`). + @inlinable + public func isDisjoint(with other: SortedSet) -> Bool { + var it1 = self.makeIterator() + var it2 = other.makeIterator() + + var e1 = it1.next() + var e2 = it2.next() + + // While both sequences have a value, consume the smallest element. + while let el1 = e1, let el2 = e2 { + if el1 < el2 { + e1 = it1.next() + } else if el1 == el2 { + return false + } else { + // el1 > el2 + e2 = it2.next() + } + } + + return true + } +} diff --git a/Sources/SortedCollections/SortedSet/SortedSet+SubSequence.swift b/Sources/SortedCollections/SortedSet/SortedSet+SubSequence.swift new file mode 100644 index 000000000..32995c5b5 --- /dev/null +++ b/Sources/SortedCollections/SortedSet/SortedSet+SubSequence.swift @@ -0,0 +1,314 @@ +//===----------------------------------------------------------------------===// +// +// 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 SortedSet { + public struct SubSequence { + @usableFromInline + internal typealias _TreeSubSequence = _Tree.SubSequence + + @usableFromInline + internal let _subSequence: _TreeSubSequence + + @inlinable + @inline(__always) + internal init(_ _subSequence: _TreeSubSequence) { + self._subSequence = _subSequence + } + + /// The underlying collection of the subsequence. + @inlinable + @inline(__always) + internal var base: SortedSet { SortedSet(_rootedAt: _subSequence.base) } + } +} + +extension SortedSet.SubSequence: Sequence { + public typealias Element = SortedSet.Element + + + public struct Iterator: IteratorProtocol { + @usableFromInline + internal var _iterator: _TreeSubSequence.Iterator + + @inlinable + @inline(__always) + internal init(_ _iterator: _TreeSubSequence.Iterator) { + self._iterator = _iterator + } + + /// Advances to the next element and returns it, or nil if no next element exists. + /// + /// - Returns: The next element in the underlying sequence, if a next element exists; + /// otherwise, `nil`. + /// - Complexity: O(1) amortized over the entire sequence. + @inlinable + @inline(__always) + public mutating func next() -> Element? { + _iterator.next()?.key + } + } + + /// Returns an iterator over the elements of the subsequence. + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + @inline(__always) + public __consuming func makeIterator() -> Iterator { + Iterator(_subSequence.makeIterator()) + } +} + +extension SortedSet.SubSequence: BidirectionalCollection { + public typealias Index = SortedSet.Index + public typealias SubSequence = Self + + /// The position of the first element in a nonempty subsequence. + /// + /// If the collection is empty, `startIndex` is equal to `endIndex`. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var startIndex: Index { Index(_subSequence.startIndex) } + + /// The subsequence's "past the end" position---that is, the position one + /// greater than the last valid subscript argument. + /// + /// If the collection is empty, `endIndex` is equal to `startIndex`. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var endIndex: Index { Index(_subSequence.endIndex) } + + /// The number of elements in the subsequence. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public var count: Int { _subSequence.count } + + /// Returns the distance between two indices. + /// + /// - 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. + /// - Complexity: O(1) + @inlinable + @inline(__always) + public func distance(from start: Index, to end: Index) -> Int { + start._index.ensureValid(forTree: _subSequence.base) + end._index.ensureValid(forTree: _subSequence.base) + return _subSequence.distance(from: start._index, to: end._index) + } + + + /// Returns the position immediately after the given index. + /// + /// - Parameter i: A valid index of the collection. `i` must be less than `endIndex`. + /// - Returns: The index value immediately after `i`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + @inline(__always) + public func index(after i: Index) -> Index { + i._index.ensureValid(forTree: _subSequence.base) + return Index(_subSequence.index(after: i._index)) + } + + /// Replaces the given index with its successor. + /// + /// - Parameter i: A valid index of the collection. `i` must be less than `endIndex`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + @inline(__always) + public func formIndex(after i: inout Index) { + i._index.ensureValid(forTree: _subSequence.base) + return _subSequence.formIndex(after: &i._index) + } + + + /// Returns the position immediately before the given index. + /// + /// - Parameter i: A valid index of the collection. `i` must be greater + /// than `startIndex`. + /// - Returns: The index value immediately before `i`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + @inline(__always) + public func index(before i: Index) -> Index { + i._index.ensureValid(forTree: _subSequence.base) + return Index(_subSequence.index(before: i._index)) + } + + /// Replaces the given index with its predecessor. + /// + /// - Parameter i: A valid index of the collection. `i` must be greater + /// than `startIndex`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + @inline(__always) + public func formIndex(before i: inout Index) { + i._index.ensureValid(forTree: _subSequence.base) + _subSequence.formIndex(before: &i._index) + } + + + /// Returns an index that is the specified distance from the given index. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - Returns: An index offset by `distance` from the index `i`. If `distance` + /// is positive, this is the same value as the result of `distance` calls to + /// `index(after:)`. If `distance` is negative, this is the same value as the + /// result of `abs(distance)` calls to `index(before:)`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + @inline(__always) + public func index(_ i: Index, offsetBy distance: Int) -> Index { + i._index.ensureValid(forTree: _subSequence.base) + return Index(_subSequence.index(i._index, offsetBy: distance)) + } + + /// Offsets the given index by the specified distance. + /// + /// The value passed as distance must not offset i beyond the bounds of the collection. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + @inline(__always) + internal func formIndex(_ i: inout Index, offsetBy distance: Int) { + i._index.ensureValid(forTree: _subSequence.base) + _subSequence.formIndex(&i._index, offsetBy: distance) + } + + + /// Returns an index that is the specified distance from the given index, unless that distance is beyond + /// a given limiting index. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - limit: A valid index of the collection to use as a limit. If `distance > 0`, a limit that is less + /// than `i` has no effect. Likewise, if `distance < 0`, a limit that is greater than `i` has + /// no effect. + /// - Returns: An index offset by `distance` from the index `i`, unless that index would be + /// beyond `limit` in the direction of movement. In that case, the method returns `nil`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + @inline(__always) + public func index(_ i: Index, offsetBy distance: Int, limitedBy limit: Index) -> Index? { + i._index.ensureValid(forTree: _subSequence.base) + limit._index.ensureValid(forTree: _subSequence.base) + + if let i = _subSequence.index(i._index, offsetBy: distance, limitedBy: limit._index) { + return Index(i) + } else { + return nil + } + } + + /// Offsets the given index by the specified distance, or so that it equals the given limiting index. + /// + /// - Parameters: + /// - i: A valid index of the collection. + /// - distance: The distance to offset `i`. + /// - limit: A valid index of the collection to use as a limit. If `distance > 0`, a limit that is less + /// than `i` has no effect. Likewise, if `distance < 0`, a limit that is greater than `i` has + /// no effect. + /// - Returns: `true` if `i` has been offset by exactly `distance` steps without going beyond + /// `limit`; otherwise, `false`. When the return value is `false`, the value of `i` is + /// equal to `limit`. + /// - Complexity: O(log(`self.count`)) in the worst-case. + @inlinable + @inline(__always) + internal func formIndex(_ i: inout Index, offsetBy distance: Int, limitedBy limit: Self.Index) -> Bool { + i._index.ensureValid(forTree: _subSequence.base) + limit._index.ensureValid(forTree: _subSequence.base) + return _subSequence.formIndex(&i._index, offsetBy: distance, limitedBy: limit._index) + } + + @inlinable + @inline(__always) + public subscript(position: Index) -> Element { + position._index.ensureValid(forTree: _subSequence.base) + return _subSequence[position._index].key + } + + @inlinable + public subscript(bounds: Range) -> SubSequence { + bounds.lowerBound._index.ensureValid(forTree: _subSequence.base) + bounds.upperBound._index.ensureValid(forTree: _subSequence.base) + + let bound = bounds.lowerBound._index..) { + _subSequence._failEarlyRangeCheck( + index._index, + bounds: bounds.lowerBound._index.., bounds: Range) { + _subSequence._failEarlyRangeCheck( + range.lowerBound._index.. Bool { + if lhs.count != rhs.count { return false } + for (k1, k2) in zip(lhs, rhs) { + if k1 != k2 { + return false + } + } + return true + } +} + +extension SortedSet.SubSequence: Hashable where Element: Hashable { + /// Hashes the essential components of this value by feeding them + /// into the given hasher. + /// - Parameter hasher: The hasher to use when combining + /// the components of this instance. + /// - Complexity: O(`self.count`) + @inlinable + public func hash(into hasher: inout Hasher) { + hasher.combine(self.count) + for element in self { + hasher.combine(element) + } + } +} diff --git a/Sources/SortedCollections/SortedSet/SortedSet+Subscripts.swift b/Sources/SortedCollections/SortedSet/SortedSet+Subscripts.swift new file mode 100644 index 000000000..a04d3771f --- /dev/null +++ b/Sources/SortedCollections/SortedSet/SortedSet+Subscripts.swift @@ -0,0 +1,36 @@ +//===----------------------------------------------------------------------===// +// +// 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 SortedSet { + @inlinable + public subscript(position: Index) -> Element { + position._index.ensureValid(forTree: self._root) + return self._root[position._index].key + } + + /// Returns a sequence of elements in the collection bounded by the provided + /// range. + /// + /// This is particularly useful when applied with a bound corresponding to some + /// group of elements. + /// + /// let students: SortedSet = ... + /// students["A"..<"B"] // Sequence of students with names beginning with "A" + /// + /// - Complexity: O(log(`self.count`)) + @inlinable + public subscript(range: Range) -> SubSequence { + let start = _root.startIndex(forKey: range.lowerBound) + let end = _root.startIndex(forKey: range.upperBound) + let range = _Tree.SubSequence(base: _root, bounds: start.. Index? { + if let index = self._root.findAnyIndex(forKey: element) { + return Index(index) + } else { + return nil + } + } + + /// The position of an element within a sorted set + public struct Index { + @usableFromInline + internal var _index: _Tree.Index + + @inlinable + @inline(__always) + internal init(_ _index: _Tree.Index) { + self._index = _index + } + } +} + +// MARK: Equatable +extension SortedSet.Index: Equatable { + @inlinable + public static func ==(lhs: SortedSet.Index, rhs: SortedSet.Index) -> Bool { + lhs._index.ensureValid(with: rhs._index) + return lhs._index == rhs._index + } +} + +// MARK: Comparable +extension SortedSet.Index: Comparable { + @inlinable + public static func <(lhs: SortedSet.Index, rhs: SortedSet.Index) -> Bool { + lhs._index.ensureValid(with: rhs._index) + return lhs._index < rhs._index + } +} diff --git a/Sources/SortedCollections/SortedSet/SortedSet.swift b/Sources/SortedCollections/SortedSet/SortedSet.swift new file mode 100644 index 000000000..15a5bde16 --- /dev/null +++ b/Sources/SortedCollections/SortedSet/SortedSet.swift @@ -0,0 +1,37 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +/// A collection which maintains unique members in ascending sorted order. +public struct SortedSet { + @usableFromInline + internal typealias _Tree = _BTree + + @usableFromInline + internal var _root: _Tree + + //// Creates an empty set. + /// + /// This initializer is equivalent to initializing with an empty array + /// literal. + /// + /// - Complexity: O(1) + @inlinable + @inline(__always) + public init() { + self._root = _Tree() + } + + /// Creates a set rooted at a given B-Tree. + @inlinable + internal init(_rootedAt tree: _Tree) { + self._root = tree + } +} diff --git a/Sources/SortedCollections/Utilities/Assertions.swift b/Sources/SortedCollections/Utilities/Assertions.swift new file mode 100644 index 000000000..fda167d42 --- /dev/null +++ b/Sources/SortedCollections/Utilities/Assertions.swift @@ -0,0 +1,12 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + + diff --git a/Tests/SortedCollectionsTests/BTree/BTree Tests.swift b/Tests/SortedCollectionsTests/BTree/BTree Tests.swift new file mode 100644 index 000000000..f9ec01733 --- /dev/null +++ b/Tests/SortedCollectionsTests/BTree/BTree Tests.swift @@ -0,0 +1,173 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +import CollectionsTestSupport +@_spi(Testing) @testable import SortedCollections + +func btreeOfSize( + _ size: Int, + _ body: (inout _BTree, [(key: Int, value: Int)]) throws -> Void +) rethrows { + var tree = _BTree(capacity: 2) + var keyValues = [(key: Int, value: Int)]() + for i in 0..() + for (key, value) in kvs { + tree.updateAnyValue(value, forKey: key) + } + } +} diff --git a/Tests/SortedCollectionsTests/BTree/BTree+Deletion Tests.swift b/Tests/SortedCollectionsTests/BTree/BTree+Deletion Tests.swift new file mode 100644 index 000000000..12774de2c --- /dev/null +++ b/Tests/SortedCollectionsTests/BTree/BTree+Deletion Tests.swift @@ -0,0 +1,31 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +import CollectionsTestSupport +@_spi(Testing) @testable import SortedCollections + +final class NodeDeletionTests: CollectionTestCase { + func test_singleDeletion() { + withEvery("size", in: [1, 2, 4, 8, 16, 32, 64, 128]) { size in + withEvery("key", in: 0...Builder(capacity: 4) + + for i in 0.. _Node { + let kvs = self.keys.map { (key: $0, value: $0 * 2) } + return _Node( + _keyValuePairs: kvs, + children: children?.map({ $0.toNode(ofCapacity: capacity) }), + capacity: capacity + ) + } + + func toBTree(ofCapacity capacity: Int) -> _BTree { + return _BTree(rootedAt: self.toNode(ofCapacity: capacity), internalCapacity: capacity) + } + + func matches(_ btree: _BTree) -> Bool { + return self.matches(btree.root) + } + + func matches(_ node: _Node) -> Bool { + return node.read { handle in + if self.keys.count != handle.elementCount { return false } + if (self.children == nil) != handle.isLeaf { return false } + + if let children = self.children { + for (i, child) in children.enumerated() { + if !child.matches(handle[childAt: i]) { + return false + } + } + } + + for (i, key) in self.keys.enumerated() { + if handle[keyAt: i] != key { + return false + } + } + + return true + } + } +} + +@resultBuilder +struct NodeTemplateBuilder { + static func buildBlock(_ components: Any...) -> NodeTemplate { + var keys = [Int]() + var children = [NodeTemplate]() + for c in components { + switch c { + case let c as Int: + keys.append(c) + case let c as NodeTemplate: + children.append(c) + default: + preconditionFailure("NodeTemplate child must be either key or child.") + } + } + precondition(children.count == 0 || children.count == keys.count + 1, + "NodeTemplate must be either leaf or internal node.") + return NodeTemplate(keys: keys, children: children.isEmpty ? nil : children) + } +} + +func tree(@NodeTemplateBuilder _ builder: () -> NodeTemplate) -> NodeTemplate { + return builder() +} diff --git a/Tests/SortedCollectionsTests/BTree/Node+Balancing Tests.swift b/Tests/SortedCollectionsTests/BTree/Node+Balancing Tests.swift new file mode 100644 index 000000000..f108148b8 --- /dev/null +++ b/Tests/SortedCollectionsTests/BTree/Node+Balancing Tests.swift @@ -0,0 +1,236 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +import CollectionsTestSupport +@_spi(Testing) @testable import SortedCollections + +final class NodeBalancingTests: CollectionTestCase { + func test_collapseAtSlot() { + let t = tree { + tree { + tree { 0; 1 } + 2 + tree { 3; 4 } + 5 + tree { 6; 7 } + } + 8 + tree { + tree { 9; 10 } + 11 + tree { 12; 13 } + 14 + tree { 15; 16 } + } + 17 + tree { + tree { 18; 19 } + 20 + tree { 21; 22 } + 23 + tree { 24; 25 } + } + } + + var btree = t.toBTree(ofCapacity: 2) + print(btree.debugDescription) + btree.removeAnyElement(forKey: 0) + btree.removeAnyElement(forKey: 1) + btree.removeAnyElement(forKey: 2) + btree.removeAnyElement(forKey: 3) + btree.removeAnyElement(forKey: 4) + btree.removeAnyElement(forKey: 5) + btree.removeAnyElement(forKey: 6) + btree.removeAnyElement(forKey: 18) + print(btree.debugDescription) +// print(SortedDictionary(_rootedAt: btree)) + } + + // MARK: Right Rotation + func test_internalRightRotation() { + let t = tree { + tree { + tree { 1; 2 } + 3 + tree { 4 ; 5 } + 6 + tree { 7; 8 } + } + 9 + tree { + tree { 10; 11 } + } + } + + var node = t.toNode(ofCapacity: 2) + node.update { $0.rotateRight(atSlot: 0) } + + expectTrue( + tree { + tree { + tree { 1; 2 } + 3 + tree { 4; 5 } + } + 6 + tree { + tree { 7; 8 } + 9 + tree { 10; 11 } + } + }.matches(node) + ) + } + + func test_leafRightRotation() { + let t = tree { + tree { 0; 1 } + 2 + tree { 3 } + } + + var node = t.toNode(ofCapacity: 2) + node.update { $0.rotateRight(atSlot: 0) } + + expectTrue( + tree { + tree { 0 } + 1 + tree { 2; 3 } + }.matches(node) + ) + } + + // MARK: Left Rotation + func test_internalLeftRotation() { + let t = tree { + tree { + tree { 1; 2 } + } + 3 + tree { + tree { 4 ; 5 } + 6 + tree { 7; 8 } + 9 + tree { 10; 11 } + } + } + + var node = t.toNode(ofCapacity: 2) + node.update { $0.rotateLeft(atSlot: 0) } + + expectTrue( + tree { + tree { + tree { 1; 2 } + 3 + tree { 4; 5 } + } + 6 + tree { + tree { 7; 8 } + 9 + tree { 10; 11 } + } + }.matches(node) + ) + } + + func test_leafLeftRotation() { + let t = tree { + tree { 0 } + 1 + tree { 2; 3 } + } + + var node = t.toNode(ofCapacity: 2) + node.update { $0.rotateLeft(atSlot: 0) } + + expectTrue( + tree { + tree { 0; 1 } + 2 + tree { 3 } + }.matches(node) + ) + } + + func test_emptyLeafLeftRotation() { + let t = tree { + tree { } + 1 + tree { 2; 3 } + } + + var node = t.toNode(ofCapacity: 2) + node.update { $0.rotateLeft(atSlot: 0) } + + expectTrue( + tree { + tree { 1 } + 2 + tree { 3 } + }.matches(node) + ) + } + + // MARK: Collapse + func test_internalCollapse() { + let t = tree { + tree { + tree { 1; 2 } + } + 3 + tree { + tree { 4 } + 5 + tree { 6; 7 } + } + } + + var node = t.toNode(ofCapacity: 2) + node.update { $0.collapse(atSlot: 0) } + + expectTrue( + tree { + tree { + tree { 1; 2 } + 3 + tree { 4 } + 5 + tree { 6; 7 } + } + }.matches(node) + ) + } + + func test_leafCollapse() { + let t = tree { + tree { } + 2 + tree { 3 } + 4 + tree { 5; 6 } + } + + var node = t.toNode(ofCapacity: 2) + node.update { $0.collapse(atSlot: 0) } + + expectTrue( + tree { + tree { 2; 3 } + 4 + tree { 5; 6 } + }.matches(node) + ) + } +} diff --git a/Tests/SortedCollectionsTests/BTree/Node+Insertion Tests.swift b/Tests/SortedCollectionsTests/BTree/Node+Insertion Tests.swift new file mode 100644 index 000000000..43227b9b6 --- /dev/null +++ b/Tests/SortedCollectionsTests/BTree/Node+Insertion Tests.swift @@ -0,0 +1,545 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +import CollectionsTestSupport +@_spi(Testing) @testable import SortedCollections + +func expectInsertionInTree( + capacity: Int, + tree: NodeTemplate, + inserting key: Int, + toEqual refTree: NodeTemplate) { + var btree = tree.toBTree(ofCapacity: capacity) + + btree.updateAnyValue(key * 2, forKey: key) + + let refMatches = refTree.matches(btree) + if !refMatches { + print("Expected: ") + print(refTree.toBTree(ofCapacity: capacity)) + print("Instead got: ") + print(btree) + } + expectTrue(refMatches) +} + +final class NodeInsertionTests: CollectionTestCase { + // MARK: Median Leaf Node Insertion + func test_medianLeafInsertion() { + expectInsertionInTree( + capacity: 2, + tree: tree { 0; 2 }, + inserting: 1, + toEqual: tree { + tree { 0 } + 1 + tree { 2 } + } + ) + + expectInsertionInTree( + capacity: 4, + tree: tree { 0; 1; 3; 4 }, + inserting: 2, + toEqual: tree { + tree { 0; 1 } + 2 + tree { 3; 4 } + } + ) + + expectInsertionInTree( + capacity: 5, + tree: tree { 0; 1; 3; 4; 5 }, + inserting: 2, + toEqual: tree { + tree { 0; 1 } + 2 + tree { 3; 4; 5 } + } + ) + } + + // MARK: Median Internal Node Insertion + func test_medianInternalInsertion() { + expectInsertionInTree( + capacity: 2, + tree: tree { + tree { 0; 1 } + 2 + tree { 3; 5 } + 6 + tree { 7; 8 } + }, + inserting: 4, + toEqual: tree { + tree { + tree { 0; 1 } + 2 + tree { 3 } + } + 4 + tree { + tree { 5 } + 6 + tree { 7; 8 } + } + } + ) + + expectInsertionInTree( + capacity: 3, + tree: tree { + tree { 0; 1; 2 } + 3 + tree { 4; 6; 7 } + 8 + tree { 9; 10; 11 } + 12 + tree { 13; 14; 15 } + }, + inserting: 5, + toEqual: tree { + tree { + tree { 0; 1; 2 } + 3 + tree { 4 } + } + 5 + tree { + tree { 6; 7 } + 8 + tree { 9; 10; 11 } + 12 + tree { 13; 14; 15 } + } + } + ) + + expectInsertionInTree( + capacity: 4, + tree: tree { + tree { 0; 1; 2; 4 } + 5 + tree { 6; 7; 8; 9 } + 10 + tree { 11; 12; 14; 15 } + 16 + tree { 17; 18; 19; 20 } + 21 + tree { 22; 23; 24; 25 } + }, + inserting: 13, + toEqual: tree { + tree { + tree { 0; 1; 2; 4 } + 5 + tree { 6; 7; 8; 9 } + 10 + tree { 11; 12 } + } + 13 + tree { + tree { 14; 15 } + 16 + tree { 17; 18; 19; 20 } + 21 + tree { 22; 23; 24; 25 } + } + } + ) + + expectInsertionInTree( + capacity: 5, + tree: tree { + tree { 0; 1; 2; 4; 5 } + 6 + tree { 7; 8; 9; 10; 11 } + 12 + tree { 13; 14; 16; 17; 18 } + 19 + tree { 20; 21; 22; 23; 24 } + 25 + tree { 26; 27; 28; 29; 30 } + 31 + tree { 32; 33; 34; 35; 36 } + }, + inserting: 15, + toEqual: tree { + tree { + tree { 0; 1; 2; 4; 5 } + 6 + tree { 7; 8; 9; 10; 11 } + 12 + tree { 13; 14 } + } + 15 + tree { + tree { 16; 17; 18 } + 19 + tree { 20; 21; 22; 23; 24 } + 25 + tree { 26; 27; 28; 29; 30 } + 31 + tree { 32; 33; 34; 35; 36 } + } + } + ) + } + + // MARK: Right Leaf Insertion + func test_rightLeafInsertion() { + expectInsertionInTree( + capacity: 2, + tree: tree { 1; 2 }, + inserting: 3, + toEqual: tree { + tree { 1 } + 2 + tree { 3 } + } + ) + + expectInsertionInTree( + capacity: 3, + tree: tree { 1; 2; 4 }, + inserting: 3, + toEqual: tree { + tree { 1 } + 2 + tree { 3; 4 } + } + ) + + expectInsertionInTree( + capacity: 4, + tree: tree { 0; 1; 2; 4 }, + inserting: 3, + toEqual: tree { + tree { 0; 1 } + 2 + tree { 3; 4 } + } + ) + + expectInsertionInTree( + capacity: 5, + tree: tree { 0; 1; 2; 3; 5 }, + inserting: 4, + toEqual: tree { + tree { 0; 1 } + 2 + tree { 3; 4; 5 } + } + ) + } + + // MARK: Right Internal Node Insertion + func test_rightInternalNodeInsertion() { + expectInsertionInTree( + capacity: 2, + tree: tree { + tree { 0; 1 } + 2 + tree { 3; 4 } + 5 + tree { 6; 7 } + }, + inserting: 8, + toEqual: tree { + tree { + tree { 0; 1 } + 2 + tree { 3; 4 } + } + 5 + tree { + tree { 6 } + 7 + tree { 8 } + } + } + ) + + expectInsertionInTree( + capacity: 3, + tree: tree { + tree { 0; 1; 2 } + 3 + tree { 4; 5; 6 } + 7 + tree { 8; 9; 10 } + 11 + tree { 12; 13; 15 } + }, + inserting: 14, + toEqual: tree { + tree { + tree { 0; 1; 2 } + 3 + tree { 4; 5; 6 } + } + 7 + tree { + tree { 8; 9; 10 } + 11 + tree { 12 } + 13 + tree { 14; 15 } + } + } + ) + + expectInsertionInTree( + capacity: 4, + tree: tree { + tree { 0; 1; 2; 4 } + 5 + tree { 6; 7; 8; 9 } + 10 + tree { 11; 12; 13; 14 } + 15 + tree { 16; 17; 18; 19 } + 20 + tree { 21; 22; 23; 25 } + }, + inserting: 24, + toEqual: tree { + tree { + tree { 0; 1; 2; 4 } + 5 + tree { 6; 7; 8; 9 } + 10 + tree { 11; 12; 13; 14 } + } + 15 + tree { + tree { 16; 17; 18; 19 } + 20 + tree { 21; 22 } + 23 + tree { 24; 25 } + } + } + ) + + expectInsertionInTree( + capacity: 5, + tree: tree { + tree { 0; 1; 2; 4; 5 } + 6 + tree { 7; 8; 9; 10; 11 } + 12 + tree { 13; 14; 15; 16; 17 } + 18 + tree { 19; 20; 21; 22; 23 } + 24 + tree { 25; 26; 27; 28; 29 } + 30 + tree { 31; 32; 33; 34; 36 } + }, + inserting: 35, + toEqual: tree { + tree { + tree { 0; 1; 2; 4; 5 } + 6 + tree { 7; 8; 9; 10; 11 } + 12 + tree { 13; 14; 15; 16; 17 } + } + 18 + tree { + tree { 19; 20; 21; 22; 23 } + 24 + tree { 25; 26; 27; 28; 29 } + 30 + tree { 31; 32 } + 33 + tree { 34; 35; 36 } + } + } + ) + } + + // MARK: Left Leaf Insertion + func test_leftLeafInsertion() { + expectInsertionInTree( + capacity: 2, + tree: tree { 1; 2 }, + inserting: 0, + toEqual: tree { + tree { 0 } + 1 + tree { 2 } + } + ) + + expectInsertionInTree( + capacity: 3, + tree: tree { 1; 2; 3 }, + inserting: 0, + toEqual: tree { + tree { 0; 1 } + 2 + tree { 3 } + } + ) + + expectInsertionInTree( + capacity: 4, + tree: tree { 0; 2; 3; 4 }, + inserting: 1, + toEqual: tree { + tree { 0; 1 } + 2 + tree { 3; 4 } + } + ) + + expectInsertionInTree( + capacity: 5, + tree: tree { 0; 2; 3; 4; 5 }, + inserting: 1, + toEqual: tree { + tree { 0; 1; 2 } + 3 + tree { 4; 5 } + } + ) + } + + // MARK: Left Internal Node Insertion + func test_leftInternalNodeInsertion() { + expectInsertionInTree( + capacity: 2, + tree: tree { + tree { 1; 2 } + 3 + tree { 4; 5 } + 6 + tree { 7; 8 } + }, + inserting: 0, + toEqual: tree { + tree { + tree { 0 } + 1 + tree { 2 } + } + 3 + tree { + tree { 4; 5 } + 6 + tree { 7; 8 } + } + } + ) + + expectInsertionInTree( + capacity: 3, + tree: tree { + tree { 1; 2; 3 } + 4 + tree { 5; 6; 7 } + 8 + tree { 9; 10; 11 } + 12 + tree { 13; 14; 15 } + }, + inserting: 0, + toEqual: tree { + tree { + tree { 0; 1 } + 2 + tree { 3 } + 4 + tree { 5; 6; 7 } + } + 8 + tree { + tree { 9; 10; 11 } + 12 + tree { 13; 14; 15 } + } + } + ) + + expectInsertionInTree( + capacity: 4, + tree: tree { + tree { 0; 2; 3; 4 } + 5 + tree { 6; 7; 8; 9 } + 10 + tree { 11; 12; 13; 14 } + 15 + tree { 16; 17; 18; 19 } + 20 + tree { 21; 22; 23; 24 } + }, + inserting: 1, + toEqual: tree { + tree { + tree { 0; 1 } + 2 + tree { 3; 4 } + 5 + tree { 6; 7; 8; 9 } + } + 10 + tree { + tree { 11; 12; 13; 14 } + 15 + tree { 16; 17; 18; 19 } + 20 + tree { 21; 22; 23; 24 } + } + } + ) + + expectInsertionInTree( + capacity: 5, + tree: tree { + tree { 0; 2; 3; 4; 5 } + 6 + tree { 7; 8; 9; 10; 11 } + 12 + tree { 13; 14; 15; 16; 17 } + 18 + tree { 19; 20; 21; 22; 23 } + 24 + tree { 25; 26; 27; 28; 29 } + 30 + tree { 31; 32; 33; 34; 35 } + }, + inserting: 1, + toEqual: tree { + tree { + tree { 0; 1; 2 } + 3 + tree { 4; 5 } + 6 + tree { 7; 8; 9; 10; 11 } + 12 + tree { 13; 14; 15; 16; 17 } + } + 18 + tree { + tree { 19; 20; 21; 22; 23 } + 24 + tree { 25; 26; 27; 28; 29 } + 30 + tree { 31; 32; 33; 34; 35 } + } + } + ) + } +} diff --git a/Tests/SortedCollectionsTests/BTree/Node+Join Tests.swift b/Tests/SortedCollectionsTests/BTree/Node+Join Tests.swift new file mode 100644 index 000000000..390272fff --- /dev/null +++ b/Tests/SortedCollectionsTests/BTree/Node+Join Tests.swift @@ -0,0 +1,106 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +import CollectionsTestSupport +@_spi(Testing) @testable import SortedCollections + +func expectNodeJoin( + capacity: Int, + tree1: NodeTemplate, + seperator: Int, + tree2: NodeTemplate, + toEqual refTree: NodeTemplate +) { + var tree1 = tree1.toNode(ofCapacity: capacity) + var tree2 = tree2.toNode(ofCapacity: capacity) + + let newTree = _Node.join( + &tree1, + with: &tree2, + seperatedBy: (seperator, -seperator), + capacity: capacity + ) + + expectTrue(refTree.matches(newTree)) + _BTree(rootedAt: newTree, internalCapacity: capacity).checkInvariants() +} + + +final class NodeJoinTests: CollectionTestCase { + func test_joinSimple() { + expectNodeJoin( + capacity: 5, + tree1: tree { + 0 + }, + seperator: 1, + tree2: tree { + 2; 3; 4 + }, + toEqual: tree { 0; 1; 2; 3; 4 } + ) + + expectNodeJoin( + capacity: 5, + tree1: tree { + 0 + }, + seperator: 1, + tree2: tree { + 2; 3; 4; 5 + }, + toEqual: tree { + tree { 0; 1; 2 } + 3 + tree { 4; 5 } + } + ) + } + + func test_joinMedian() { + expectNodeJoin( + capacity: 2, + tree1: tree { + tree { 0 } + 1 + tree { 2 } + 2 + tree { 3 } + }, + seperator: 4, + tree2: tree { + tree { 5 } + 6 + tree { 7 } + 8 + tree { 9 } + }, + toEqual: tree { + tree { + tree { 0 } + 1 + tree { 2 } + 2 + tree { 3 } + } + 4 + tree { + tree { 5 } + 6 + tree { 7 } + 8 + tree { 9 } + } + } + ) + } +} + diff --git a/Tests/SortedCollectionsTests/BTree/NodeTests.swift b/Tests/SortedCollectionsTests/BTree/NodeTests.swift new file mode 100644 index 000000000..722f558b4 --- /dev/null +++ b/Tests/SortedCollectionsTests/BTree/NodeTests.swift @@ -0,0 +1,143 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +import CollectionsTestSupport +@_spi(Testing) @testable import SortedCollections + +func nodeFromKeys(_ keys: [Int], capacity: Int) -> _Node { + let kvPairs = keys.map { (key: $0, value: $0 * 2) } + return _Node(_keyValuePairs: kvPairs, capacity: capacity) +} + +func insertSortedValue(_ value: Int, into array: inout [Int]) { + var insertionIndex = 0 + while insertionIndex < array.count { + if array[insertionIndex] > value { + break + } + insertionIndex += 1 + } + array.insert(value, at: insertionIndex) +} + +func findFirstIndexOf(_ value: Int, in array: [Int]) -> Int { + var index = 0 + while index < array.count { + if array[index] >= value { + break + } + index += 1 + } + return index +} + +func findLastIndexOf(_ value: Int, in array: [Int]) -> Int { + var index = 0 + while index < array.count { + if array[index] > value { + break + } + index += 1 + } + return index +} + +/// Generates all shifts of a duplicate run in a node of capacity N. +/// - Parameters: +/// - capacity: Total capacity of node. +/// - keys: The number filled keys in the node. +/// - duplicates: The number of duplicates. Must be greater than or equal to 1 +/// - Returns: The duplicated key. +func withEveryNode( + ofCapacity capacity: Int, + keys: Int, + duplicates: Int, + _ body: (_Node, [Int], Int) throws -> Void +) rethrows { + let possibleShifts = keys - duplicates + 1 + try withEvery("shift", in: 0...Splinter? = node.update { handle in + let index = handle.endSlot(forKey: newKey) + return handle.insertElement((newKey, newKey * 2), withRightChild: nil, atSlot: index) + } + insertSortedValue(newKey, into: &array) + + expectNil(splinter) + node.read { handle in + let keys = UnsafeBufferPointer(start: handle.keys, count: handle.elementCount) + expectEqualElements(keys, array) + expectEqual(handle.subtreeCount, count + 1) + } + } + } + } + } + + func test_firstIndexOfDuplicates() { + withEvery("capacity", in: 2..<10) { capacity in + withEvery("keys", in: 0...capacity) { keys in + withEvery("duplicates", in: 0...keys) { duplicates in + withEveryNode(ofCapacity: capacity, keys: keys, duplicates: duplicates) { node, array, duplicatedKey in + node.read { handle in + expectEqual( + handle.startSlot(forKey: duplicatedKey), + findFirstIndexOf(duplicatedKey, in: array) + ) + } + } + } + } + } + } + + func test_lastIndexOfDuplicates() { + withEvery("capacity", in: 2..<10) { capacity in + withEvery("keys", in: 0...capacity) { keys in + withEvery("duplicates", in: 0...keys) { duplicates in + withEveryNode(ofCapacity: capacity, keys: keys, duplicates: duplicates) { node, array, duplicatedKey in + node.read { handle in + expectEqual( + handle.endSlot(forKey: duplicatedKey), + findLastIndexOf(duplicatedKey, in: array) + ) + } + } + } + } + } + } +} diff --git a/Tests/SortedCollectionsTests/SortedDictionary/SortedDictionary Tests.swift b/Tests/SortedCollectionsTests/SortedDictionary/SortedDictionary Tests.swift new file mode 100644 index 000000000..10bd69ae1 --- /dev/null +++ b/Tests/SortedCollectionsTests/SortedDictionary/SortedDictionary Tests.swift @@ -0,0 +1,282 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +import CollectionsTestSupport +@_spi(Testing) @testable import SortedCollections + +final class SortedDictionaryTests: CollectionTestCase { + func test_empty() { + let d = SortedDictionary() + expectEqualElements(d, []) + expectEqual(d.count, 0) + } + + func test_keysWithValues_unique() { + let items: KeyValuePairs = [ + 3: "three", + 1: "one", + 0: "zero", + 2: "two", + ] + let d = SortedDictionary(keysWithValues: items) + expectEqualElements(d, [ + (key: 0, value: "zero"), + (key: 1, value: "one"), + (key: 2, value: "two"), + (key: 3, value: "three") + ]) + } + + func test_keysWithValues_bulk() { + withEvery("count", in: [0, 1, 2, 4, 8, 16, 32, 64, 128, 1024, 4096]) { count in + let kvs = (0..(keysWithValues: kvs) + expectEqual(sortedDictionary.count, count) + } + } + + func test_keysWithValues_duplicates() { + let items: KeyValuePairs = [ + 3: "three", + 1: "one", + 1: "one-1", + 0: "zero", + 3: "three-1", + 2: "two", + ] + let d = SortedDictionary(keysWithValues: items) + expectEqualElements(d, [ + (key: 0, value: "zero"), + (key: 1, value: "one-1"), + (key: 2, value: "two"), + (key: 3, value: "three-1") + ]) + } + + func test_grouping_initializer() { + let items: [String] = [ + "one", "two", "three", "four", "five", + "six", "seven", "eight", "nine", "ten" + ] + let d = SortedDictionary(grouping: items, by: { $0.count }) + expectEqualElements(d, [ + (key: 3, value: ["one", "two", "six", "ten"]), + (key: 4, value: ["four", "five", "nine"]), + (key: 5, value: ["three", "seven", "eight"]), + ]) + } + + func test_ExpressibleByDictionaryLiteral() { + let d0: SortedDictionary = [:] + expectTrue(d0.isEmpty) + + let d1: SortedDictionary = [ + 1: "one", + 2: "two", + 3: "three", + 4: "four", + ] + expectEqualElements(d1.map { $0.key }, [1, 2, 3, 4]) + expectEqualElements(d1.map { $0.value }, ["one", "two", "three", "four"]) + } + + func test_counts() { + withEvery("count", in: 0 ..< 30) { count in + withLifetimeTracking { tracker in + let (d, _) = tracker.sortedDictionary(keys: 0 ..< count) + expectEqual(d.isEmpty, count == 0) + expectEqual(d.count, count) + expectEqual(d.underestimatedCount, count) + } + } + } + + func test_bidirectionalCollection() { + withEvery("count", in: [1, 2, 4, 8, 16, 32, 64]) { count in + withLifetimeTracking { tracker in + let (d, kvs) = tracker.sortedDictionary(keys: 0 ..< count) + + checkBidirectionalCollection( + d, + expectedContents: kvs, + by: { $0.key == $1.key && $0.value == $1.value } + ) + } + } + } + + func test_orderedInsertion() { + withEvery("count", in: [0, 1, 2, 3, 4, 8, 16, 64]) { count in + var sortedDictionary: SortedDictionary = [:] + + for i in 0.. = [:] + + for i in (0..(keysWithValues: kvs) + sortedDictionary[i * 2] = -i + + var comparison = Array(kvs) + comparison.insert((key: i * 2, value: -i), at: i) + + expectEqualElements(comparison, sortedDictionary) + } + } + } + + func test_subscriptSet() { + withEvery("count", in: [1, 2, 4, 8, 16, 32, 64, 512]) { count in + var sortedDictionary: SortedDictionary = [:] + + for i in 0.., LifetimeTracked> = [:] + let fallback = tracker.instance(for: -2) + withEvery("offset", in: 0 ..< count) { offset in + withHiddenCopies(if: isShared, of: &d) { d in + let key = keys[offset] + d.modifyValue(forKey: key, default: fallback) { value in + expectEqual(value, fallback) + value = values[offset] + } + expectEqual(d.count, offset + 1) + withEvery("i", in: 0 ... offset) { i in + let v = d[keys[i]] + expectEqual(v, values[i]) + } + } + } + } + } + } + } + + func test_modifySubscriptRemoval() { + func modify(_ value: inout Int?, setTo newValue: Int?) { + value = newValue + } + + withEvery("count", in: [1, 2, 4, 8, 16, 32, 64, 512]) { count in + let kvs = (0..(keysWithValues: kvs) + + withEvery("isShared", in: [false, true]) { isShared in + withHiddenCopies(if: isShared, of: &d) { d in + modify(&d[key], setTo: nil) + var comparisonKeys = Array(0.. = [:] + + withHiddenCopies(if: isShared, of: &d) { d in + for i in 0..( + keys: Keys + ) -> ( + dictionary: SortedDictionary, LifetimeTracked>, + kvs: [(LifetimeTracked, LifetimeTracked)] + ) + where Keys.Element == Int + { + let k = Array(keys) + let keys = self.instances(for: k) + let values = self.instances(for: k.map { -($0 + 1) }) + + let kvs = Array(zip(keys, values)) + + let dictionary = SortedDictionary(keysWithValues: kvs) + + return (dictionary, kvs) + } +} diff --git a/Tests/SortedCollectionsTests/SortedSet/SortedSet Tests.swift b/Tests/SortedCollectionsTests/SortedSet/SortedSet Tests.swift new file mode 100644 index 000000000..7ee5c7b03 --- /dev/null +++ b/Tests/SortedCollectionsTests/SortedSet/SortedSet Tests.swift @@ -0,0 +1,448 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +import XCTest +@_spi(Testing) import SortedCollections +import CollectionsTestSupport + +class SortedSetTests: CollectionTestCase { + func test_init_sortedElements() { + withEvery("count", in: 0 ..< 40) { count in + let set = SortedSet(sortedElements: 0 ..< count) + expectEqual(set.count, count) + expectEqual(set.isEmpty, count == 0) + expectEqualElements(set, 0 ..< count) + for i in 0 ..< count { + expectTrue(set.contains(i)) + } + } + } + + func test_init_empty() { + let set = SortedSet() + expectEqual(set.count, 0) + expectTrue(set.isEmpty) + expectEqualElements(set, []) + } + + func test_init_self() { + withEvery("count", in: 0 ..< 40) { count in + let set = SortedSet(0 ..< count) + let copy = SortedSet(set) + expectEqualElements(copy, set) + } + } + + func test_init_set() { + withEvery("count", in: 0 ..< 40) { count in + let set = Set(0 ..< count) + let sorted = SortedSet(set) + expectEqual(sorted.count, count) + expectEqualElements(sorted, set.sorted()) + } + } + + func test_init_dictionary_keys() { + withEvery("count", in: 0 ..< 20) { count in + let dict: [Int: Int] + = .init(uniqueKeysWithValues: (0 ..< count).lazy.map { (key: $0, value: 2 * $0) }) + let sorted = SortedSet(dict.keys) + expectEqual(sorted.count, count) + expectEqualElements(sorted, dict.keys.sorted()) + } + } + + func test_firstIndexOf_lastIndexOf() { + withEvery("count", in: 0 ..< 20) { count in + let contents = Array(0 ..< count) + withEvery("dupes", in: 1 ... 3) { dupes in + let input = (0 ..< count).flatMap { repeatElement($0, count: dupes) } + let set = SortedSet(input) + withEvery("item", in: contents) { item in + expectNotNil(set.firstIndex(of: item)) { index in + expectEqual(set[index], item) + let offset = set.distance(from: set.startIndex, to: index) + expectEqual(contents[offset], item) + expectEqual(set.lastIndex(of: item), index) + } + } + expectNil(set.firstIndex(of: count)) + expectNil(set.lastIndex(of: count)) + } + } + } + + func test_CustomStringConvertible() { + let a: SortedSet = [] + expectEqual(a.description, "[]") + + let b: SortedSet = [0] + expectEqual(b.description, "[0]") + + let c: SortedSet = [0, 1, 2, 3, 4] + expectEqual(c.description, "[0, 1, 2, 3, 4]") + } + + func test_CustomDebugStringConvertible() { + let a: SortedSet = [] + expectEqual(a.debugDescription, "SortedSet([])") + + let b: SortedSet = [0] + expectEqual(b.debugDescription, "SortedSet([0])") + + let c: SortedSet = [0, 1, 2, 3, 4] + expectEqual(c.debugDescription, "SortedSet([0, 1, 2, 3, 4])") + } + + func test_customReflectable() { + do { + let set: SortedSet = [1, 2, 3] + let mirror = Mirror(reflecting: set) + expectEqual(mirror.displayStyle, .set) + expectNil(mirror.superclassMirror) + expectTrue(mirror.children.compactMap { $0.label }.isEmpty) // No label + expectEqualElements(mirror.children.map { $0.value as? Int }, set.map { $0 }) + } + } + + func test_ExpressibleByArrayLiteral() { + do { + let set: SortedSet = [] + expectEqualElements(set, [] as [Int]) + } + + do { + let set: SortedSet = [1, 2, 3] + expectEqualElements(set, 1 ... 3) + } + + do { + let set: SortedSet = [ + 1, 2, 3, 4, 5, 6, 7, 8, + 1, 2, 3, 4, 5, 6, 7, 8, + 1, 2, 3, 4, 5, 6, 7, 8, + 1, 2, 3, 4, 5, 6, 7, 8, + ] + expectEqualElements(set, 1 ... 8) + } + + do { + let set: SortedSet = [ + 1, 1, 1, 1, + 2, 2, 2, 2, + 3, 3, 3, 3, + 4, 4, 4, 4, + 5, 5, 5, 5, + 6, 6, 6, 6, + 7, 7, 7, 7, + 8, 8, 8, 8, + ] + expectEqualElements(set, 1 ... 8) + } + + do { + let set: SortedSet = [ + 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] + expectEqualElements(set, 1 ... 32) + } + } + + func test_Encodable() throws { + let s1: SortedSet = [] + let v1: MinimalEncoder.Value = .array([]) + expectEqual(try MinimalEncoder.encode(s1), v1) + + let s2: SortedSet = [0, 1, 2, 3] + let v2: MinimalEncoder.Value = .array([.int(0), .int(1), .int(2), .int(3)]) + expectEqual(try MinimalEncoder.encode(s2), v2) + + let s4 = SortedSet(0 ..< 100) + let v4: MinimalEncoder.Value = .array((0 ..< 100).map { .int($0) }) + expectEqual(try MinimalEncoder.encode(s4), v4) + } + + func test_Decodable() throws { + let s1: SortedSet = [] + let v1: MinimalEncoder.Value = .array([]) + expectEqual(try MinimalDecoder.decode(v1, as: SortedSet.self), s1) + + let s2: SortedSet = [0, 1, 2, 3] + let v2: MinimalEncoder.Value = .array([.int(0), .int(1), .int(2), .int(3)]) + expectEqual(try MinimalDecoder.decode(v2, as: SortedSet.self), s2) + + let s3 = SortedSet(0 ..< 100) + let v3: MinimalEncoder.Value = .array((0 ..< 100).map { .int($0) }) + expectEqual(try MinimalDecoder.decode(v3, as: SortedSet.self), s3) + + expectThrows(try MinimalDecoder.decode(.int(0), as: SortedSet.self)) + + let v4: MinimalEncoder.Value = .array([.int(0), .int(1), .int(0)]) + expectThrows(try MinimalDecoder.decode(v4, as: SortedSet.self)) { error in + expectNotNil(error as? DecodingError) { error in + guard case .dataCorrupted(let context) = error else { + expectFailure("Unexpected error \(error)") + return + } + expectEqual(context.debugDescription, + "Decoded elements out of order.") + } + } + } + + + func withSampleRanges( + file: StaticString = #file, + line: UInt = #line, + _ body: (Range, Range) throws -> Void + ) rethrows { + for c1 in [0, 10, 32, 64, 128, 256] { + for c2 in [0, 10, 32, 64, 128, 256] { + for overlap in Set([0, 1, c1 / 2, c1, -5]) { + let r1 = 0 ..< c1 + let r2 = c1 - overlap ..< c1 - overlap + c2 + if r1.lowerBound <= r2.lowerBound { + let e1 = context.push("range1: \(r1)", file: file, line: line) + let e2 = context.push("range2: \(r2)", file: file, line: line) + defer { + context.pop(e2) + context.pop(e1) + } + try body(r1, r2) + } else { + let e1 = context.push("range1: \(r2)", file: file, line: line) + let e2 = context.push("range2: \(r1)", file: file, line: line) + defer { + context.pop(e2) + context.pop(e1) + } + try body(r2, r1) + } + } + } + } + } + + func test_union_Self() { + withSampleRanges { r1, r2 in + let expected = Set(r1).union(r2).sorted() + + let u1 = SortedSet(r1) + let u2 = SortedSet(r2) + let actual1 = u1.union(u2) + expectEqualElements(actual1, expected) + + let actual2 = actual1.union(u2).union(u1) + expectEqualElements(actual2, expected) + } + } + + func test_formUnion_Self() { + withSampleRanges { r1, r2 in + let expected = Set(r1).union(r2).sorted() + + var res: SortedSet = [] + + let u1 = SortedSet(r1) + res.formUnion(u1) + expectEqualElements(res, r1) + + let u2 = SortedSet(r2) + res.formUnion(u2) + expectEqualElements(res, expected) + + res.formUnion(u1) + res.formUnion(u2) + expectEqualElements(res, expected) + } + } + + func test_intersection_Self() { + withSampleRanges { r1, r2 in + let expected = Set(r1).intersection(r2).sorted() + + let u1 = SortedSet(r1) + let u2 = SortedSet(r2) + let actual1 = u1.intersection(u2) + expectEqualElements(actual1, expected) + + let actual2 = actual1.intersection(u1) + expectEqualElements(actual2, expected) + } + } + + func test_formIntersection_Self() { + withSampleRanges { r1, r2 in + let expected = Set(r1).intersection(r2).sorted() + + let u1 = SortedSet(r1) + let u2 = SortedSet(r2) + var res = u1 + res.formIntersection(u2) + expectEqualElements(res, expected) + expectEqualElements(u1, r1) + + res.formIntersection(u1) + res.formIntersection(u2) + expectEqualElements(res, expected) + } + } + + + func test_symmetricDifference_Self() { + withSampleRanges { r1, r2 in + let expected = Set(r1).symmetricDifference(r2).sorted() + + let u1 = SortedSet(r1) + let u2 = SortedSet(r2) + let actual1 = u1.symmetricDifference(u2) + expectEqualElements(actual1, expected) + + let actual2 = actual1.symmetricDifference(u1).symmetricDifference(u2) + expectEqual(actual2.count, 0) + } + } + + func test_formSymmetricDifference_Self() { + withSampleRanges { r1, r2 in + let expected = Set(r1).symmetricDifference(r2).sorted() + + let u1 = SortedSet(r1) + let u2 = SortedSet(r2) + var res = u1 + res.formSymmetricDifference(u2) + expectEqualElements(res, expected) + expectEqualElements(u1, r1) + + res.formSymmetricDifference(u1) + res.formSymmetricDifference(u2) + expectEqual(res.count, 0) + } + } + + func test_subtracting_Self() { + withSampleRanges { r1, r2 in + let expected = Set(r1).subtracting(r2).sorted() + + let u1 = SortedSet(r1) + let u2 = SortedSet(r2) + let actual1 = u1.subtracting(u2) + expectEqualElements(actual1, expected) + + let actual2 = actual1.subtracting(u2) + expectEqualElements(actual2, expected) + } + } + + func test_subtract_Self() { + withSampleRanges { r1, r2 in + let expected = Set(r1).subtracting(r2).sorted() + + let u1 = SortedSet(r1) + let u2 = SortedSet(r2) + var res = u1 + res.subtract(u2) + expectEqualElements(res, expected) + expectEqualElements(u1, r1) + + res.subtract(u2) + expectEqualElements(res, expected) + } + } + + struct SampleRanges { + let unit: Int + + init(unit: Int) { + self.unit = unit + } + + var empty: Range { unit ..< unit } + + var a: Range { 0 ..< unit } + var b: Range { unit ..< 2 * unit } + var c: Range { 2 * unit ..< 3 * unit } + + var ab: Range { 0 ..< 2 * unit } + var bc: Range { unit ..< 3 * unit } + + var abc: Range { 0 ..< 3 * unit } + + var ranges: [Range] { [empty, a, b, c, ab, bc, abc] } + + func withEveryPair( + _ body: (Range, Range) throws -> Void + ) rethrows { + try withEvery("range1", in: ranges) { range1 in + try withEvery("range2", in: ranges) { range2 in + try body(range1, range2) + } + } + } + } + + func test_isSubset_Self() { + withEvery("unit", in: [1, 3, 7, 10, 20, 50]) { unit in + SampleRanges(unit: unit).withEveryPair { r1, r2 in + let expected = Set(r1).isSubset(of: r2) + let a = SortedSet(r1) + let b = SortedSet(r2) + expectEqual(a.isSubset(of: b), expected) + } + } + } + + func test_isSuperset_Self() { + withEvery("unit", in: [1, 3, 7, 10, 20, 50]) { unit in + SampleRanges(unit: unit).withEveryPair { r1, r2 in + let expected = Set(r1).isSuperset(of: r2) + let a = SortedSet(r1) + let b = SortedSet(r2) + expectEqual(a.isSuperset(of: b), expected) + } + } + } + + func test_isStrictSubset_Self() { + withEvery("unit", in: [1, 3, 7, 10, 20, 50]) { unit in + SampleRanges(unit: unit).withEveryPair { r1, r2 in + let expected = Set(r1).isStrictSubset(of: r2) + let a = SortedSet(r1) + let b = SortedSet(r2) + expectEqual(a.isStrictSubset(of: b), expected) + } + } + } + + func test_isStrictSuperset_Self() { + withEvery("unit", in: [1, 3, 7, 10, 20, 50]) { unit in + SampleRanges(unit: unit).withEveryPair { r1, r2 in + let expected = Set(r1).isStrictSuperset(of: r2) + let a = SortedSet(r1) + let b = SortedSet(r2) + expectEqual(a.isStrictSuperset(of: b), expected) + } + } + } + + func test_isDisjoint_Self() { + withEvery("unit", in: [1, 3, 7, 10, 20, 50]) { unit in + SampleRanges(unit: unit).withEveryPair { r1, r2 in + let expected = Set(r1).isDisjoint(with: r2) + let a = SortedSet(r1) + let b = SortedSet(r2) + expectEqual(a.isDisjoint(with: b), expected) + } + } + } +}