Add RangeSet and discontiguous collection operations (#28161)

This adds the RangeSet and DiscontiguousSlice types, as well as collection
operations for working with discontiguous ranges of elements. This also adds
a COWLoggingArray type to the test suite to verify that mutable collection
algorithms don't perform unexpected copy-on-write operations when mutating
slices mid-operation.

Author: natecook1000

stdlib/private/StdlibCollectionUnittest/CMakeLists.txt

@@ -12,6 +12,7 @@ add_swift_target_library(swiftStdlibCollectionUnittest ${SWIFT_STDLIB_LIBRARY_BU
   CheckRangeReplaceableSliceType.swift
   CheckSequenceInstance.swift
   CheckSequenceType.swift
+  COWLoggingArray.swift
   LoggingWrappers.swift
   MinimalCollections.swift
   RangeSelection.swift

stdlib/private/StdlibCollectionUnittest/COWLoggingArray.swift

@@ -0,0 +1,123 @@
+import StdlibUnittest
+
+fileprivate var COWLoggingArray_CopyCount = 0
+
+public func expectNoCopyOnWrite<T>(
+  _ elements: [T], 
+  _ message: @autoclosure () -> String = "",
+  stackTrace: SourceLocStack = SourceLocStack(),
+  showFrame: Bool = true,
+  file: String = #file,
+  line: UInt = #line,
+  _ body: (inout COWLoggingArray<T>) -> Void
+) {
+  let copyCountBeforeBody = COWLoggingArray_CopyCount
+  var loggingArray = COWLoggingArray(elements)
+  body(&loggingArray)
+  expectEqual(copyCountBeforeBody, COWLoggingArray_CopyCount, message(),
+    stackTrace: stackTrace.pushIf(showFrame, file: file, line: line),
+    showFrame: false)
+}
+
+public struct COWLoggingArray<Element> {
+  var storage: Storage
+  
+  class Storage {
+    var buffer: UnsafeMutableBufferPointer<Element>
+    var count: Int
+    var capacity: Int {
+      buffer.count
+    }
+    
+    init(capacity: Int) {
+      self.buffer = .allocate(capacity: capacity)
+      self.count = 0
+    }
+    
+    deinit {
+      buffer.baseAddress!.deinitialize(count: count)
+      buffer.deallocate()
+    }
+    
+    func cloned(capacity: Int? = nil) -> Storage {
+      let newCapacity = Swift.max(capacity ?? self.capacity, self.capacity)
+      let newStorage = Storage(capacity: newCapacity)
+      newStorage.buffer.baseAddress!
+        .initialize(from: buffer.baseAddress!, count: count)
+      newStorage.count = count
+      return newStorage
+    }
+  }
+  
+  mutating func _makeUnique() {
+    if !isKnownUniquelyReferenced(&storage) {
+      storage = storage.cloned()
+      COWLoggingArray_CopyCount += 1
+    }
+  }
+}
+
+extension COWLoggingArray: RandomAccessCollection, RangeReplaceableCollection,
+  MutableCollection, ExpressibleByArrayLiteral
+{
+  public var count: Int { storage.count }
+  public var startIndex: Int { 0 }
+  public var endIndex: Int { count }
+  
+  public subscript(i: Int) -> Element {
+    get {
+      storage.buffer[i]
+    }
+    set {
+      _makeUnique()
+      storage.buffer[i] = newValue
+    }
+  }
+  
+  public init() {
+    storage = Storage(capacity: 10)
+  }
+  
+  public mutating func reserveCapacity(_ n: Int) {
+    if !isKnownUniquelyReferenced(&storage) {
+      COWLoggingArray_CopyCount += 1
+      storage = storage.cloned(capacity: n)
+    } else if count < n {
+      storage = storage.cloned(capacity: n)
+    }
+  }
+  
+  public mutating func replaceSubrange<C>(_ subrange: Range<Int>, with newElements: C)
+    where C : Collection, Element == C.Element
+  {
+    _makeUnique()
+    let newCount = (count - subrange.count) + newElements.count
+    if newCount > storage.capacity {
+      storage = storage.cloned(capacity: newCount)
+    }
+    
+    let startOfSubrange = storage.buffer.baseAddress! + subrange.lowerBound
+    let endOfSubrange = startOfSubrange + subrange.count
+    let endOfNewElements = startOfSubrange + newElements.count
+    let countAfterSubrange = count - subrange.upperBound
+    
+    // clear out old elements
+    startOfSubrange.deinitialize(count: subrange.count)
+    
+    // move elements above subrange
+    endOfNewElements.moveInitialize(from: endOfSubrange, count: countAfterSubrange)
+    
+    // assign new elements
+    for (pointer, element) in zip(startOfSubrange..., newElements) {
+      pointer.initialize(to: element)
+    }
+    
+    // update count
+    storage.count = newCount
+  }
+  
+  public init(arrayLiteral elements: Element...) {
+    storage = Storage(capacity: elements.count)
+    replaceSubrange(0..<0, with: elements)
+  }
+}

stdlib/public/core/CMakeLists.txt

@@ -206,10 +206,13 @@ set(SWIFTLIB_SOURCES
   Availability.swift
   CollectionDifference.swift
   CollectionOfOne.swift
+  DiscontiguousSlice.swift
   Diffing.swift
   Mirror.swift
   PlaygroundDisplay.swift
   CommandLine.swift
+  RangeSet.swift
+  RangeSetStorage.swift
   SliceBuffer.swift
   SIMDVector.swift
   UnfoldSequence.swift

stdlib/public/core/CollectionAlgorithms.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Swift.org open source project
 //
-// Copyright (c) 2014 - 2018 Apple Inc. and the Swift project authors
+// Copyright (c) 2014 - 2020 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
@@ -209,6 +209,70 @@ extension BidirectionalCollection where Element: Equatable {
   }
 }
 
+//===----------------------------------------------------------------------===//
+// subranges(where:) / subranges(of:)
+//===----------------------------------------------------------------------===//
+
+extension Collection {
+  /// Returns the indices of all the elements that match the given predicate.
+  ///
+  /// For example, you can use this method to find all the places that a
+  /// vowel occurs in a string.
+  ///
+  ///     let str = "Fresh cheese in a breeze"
+  ///     let vowels: Set<Character> = ["a", "e", "i", "o", "u"]
+  ///     let allTheVowels = str.subranges(where: { vowels.contains($0) })
+  ///     // str[allTheVowels].count == 9
+  ///
+  /// - Parameter predicate: A closure that takes an element as its argument
+  ///   and returns a Boolean value that indicates whether the passed element
+  ///   represents a match.
+  /// - Returns: A set of the indices of the elements for which `predicate`
+  ///   returns `true`.
+  ///
+  /// - Complexity: O(*n*), where *n* is the length of the collection.
+  @available(macOS 9999, iOS 9999, tvOS 9999, watchOS 9999, *)
+  public func subranges(where predicate: (Element) throws -> Bool) rethrows
+    -> RangeSet<Index>
+  {
+    if isEmpty { return RangeSet() }
+
+    var result = RangeSet<Index>()
+    var i = startIndex
+    while i != endIndex {
+      let next = index(after: i)
+      if try predicate(self[i]) {
+        result._append(i..<next)
+      }
+      i = next
+    }
+
+    return result
+  }
+}
+
+extension Collection where Element: Equatable {
+  /// Returns the indices of all the elements that are equal to the given
+  /// element.
+  ///
+  /// For example, you can use this method to find all the places that a
+  /// particular letter occurs in a string.
+  ///
+  ///     let str = "Fresh cheese in a breeze"
+  ///     let allTheEs = str.subranges(of: "e")
+  ///     // str[allTheEs].count == 7
+  ///
+  /// - Parameter element: An element to look for in the collection.
+  /// - Returns: A set of the indices of the elements that are equal to
+  ///   `element`.
+  ///
+  /// - Complexity: O(*n*), where *n* is the length of the collection.
+  @available(macOS 9999, iOS 9999, tvOS 9999, watchOS 9999, *)
+  public func subranges(of element: Element) -> RangeSet<Index> {
+    subranges(where: { $0 == element })
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // partition(by:)
 //===----------------------------------------------------------------------===//
@@ -369,6 +433,81 @@ extension MutableCollection where Self: BidirectionalCollection {
   }
 }
 
+//===----------------------------------------------------------------------===//
+// _indexedStablePartition / _partitioningIndex
+//===----------------------------------------------------------------------===//
+
+extension MutableCollection {
+  /// Moves all elements at the indices satisfying `belongsInSecondPartition`
+  /// into a suffix of the collection, preserving their relative order, and
+  /// returns the start of the resulting suffix.
+  ///
+  /// - Complexity: O(*n* log *n*) where *n* is the number of elements.
+  /// - Precondition:
+  ///   `n == distance(from: range.lowerBound, to: range.upperBound)`
+  internal mutating func _indexedStablePartition(
+    count n: Int,
+    range: Range<Index>,
+    by belongsInSecondPartition: (Index) throws-> Bool
+  ) rethrows -> Index {
+    if n == 0 { return range.lowerBound }
+    if n == 1 {
+      return try belongsInSecondPartition(range.lowerBound)
+        ? range.lowerBound
+        : range.upperBound
+    }
+    let h = n / 2, i = index(range.lowerBound, offsetBy: h)
+    let j = try _indexedStablePartition(
+      count: h,
+      range: range.lowerBound..<i,
+      by: belongsInSecondPartition)
+    let k = try _indexedStablePartition(
+      count: n - h,
+      range: i..<range.upperBound,
+      by: belongsInSecondPartition)
+    return _rotate(in: j..<k, shiftingToStart: i)
+  }
+}
+
+//===----------------------------------------------------------------------===//
+// _partitioningIndex(where:)
+//===----------------------------------------------------------------------===//
+
+extension Collection {
+  /// Returns the index of the first element in the collection that matches
+  /// the predicate.
+  ///
+  /// The collection must already be partitioned according to the predicate.
+  /// That is, there should be an index `i` where for every element in
+  /// `collection[..<i]` the predicate is `false`, and for every element
+  /// in `collection[i...]` the predicate is `true`.
+  ///
+  /// - Parameter predicate: A predicate that partitions the collection.
+  /// - Returns: The index of the first element in the collection for which
+  ///   `predicate` returns `true`.
+  ///
+  /// - Complexity: O(log *n*), where *n* is the length of this collection if
+  ///   the collection conforms to `RandomAccessCollection`, otherwise O(*n*).
+  internal func _partitioningIndex(
+    where predicate: (Element) throws -> Bool
+  ) rethrows -> Index {
+    var n = count
+    var l = startIndex
+    
+    while n > 0 {
+      let half = n / 2
+      let mid = index(l, offsetBy: half)
+      if try predicate(self[mid]) {
+        n = half
+      } else {
+        l = index(after: mid)
+        n -= half + 1
+      }
+    }
+    return l
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // shuffled()/shuffle()
 //===----------------------------------------------------------------------===//