SortedCollection.swift

// MARK: Sorted Protocols

/// A collection with elements that are always in order according to its
/// `areInIncreasingOrder` property.
public protocol SortedCollection : Collection {
  // FIXME: unsupported generic features
  // associatedtype SubSequence : SortedCollection
  
  var areInIncreasingOrder: (Iterator.Element, Iterator.Element) -> Bool { get }

  /// Returns the insertion point for the given element.
  ///
  /// The insertion point will be immediately after any elements that compare
  /// as equal to `element` and before any elements that compare as greater
  /// than `element`.
  func stableInsertionPoint(for element: Iterator.Element) -> Index
}

/// A mutable, always-sorted collection.
public protocol MutableSortedCollection : SortedCollection {
  /// Inserts the given element and returns its new position.
  @discardableResult
  mutating func insert(_ element: Iterator.Element) -> Index
  
  /// Customization point for faster insertion.
  ///
  /// The index passed as `i` must be the correct, stable insertion point for
  /// `element`.
  mutating func _insert(_ element: Iterator.Element, atSortedIndex i: Index)
  
  /// Inserts the contents of `newElements` into the collection.
  ///
  /// `newElements` does not need to be sorted before calling
  /// `insert(contentsOf:)`.
  mutating func insert<S: Sequence>(contentsOf newElements: S)
    where S.Iterator.Element == Iterator.Element
  
  /// Inserts the contents of the previously-sorted `newElements` sequence.
  ///
  /// `newElements` must already be sorted using the same
  /// `areInIncreasingOrder` predicate as this collection.
  mutating func insert<S: Sequence>(sortedContentsOf newElements: S)
    where S.Iterator.Element == Iterator.Element
  
  /// Removes the element at `i`.
  @discardableResult
  mutating func remove(at i: Index) -> Iterator.Element
  
  /// Removes the elements in the given subrange.
  mutating func removeSubrange(_ bounds: Range<Index>)
  
  /// Removes all elements from the collection.
  mutating func removeAll(keepingCapacity: Bool)
}

// MARK: Default Implementations

extension SortedCollection {
  public func stableInsertionPoint(for element: Iterator.Element) -> Index {
    return _partitionPoint(where: { areInIncreasingOrder(element, $0) })
  }
  
  // Workaround for missing generic features.
  public func _stableInsertionPoint(for element: Iterator.Element, startingAt i: Index) -> Index {
    return _partitionPoint(where: { areInIncreasingOrder(element, $0) }, startingAt: i)
  }
  
  // Customization point for efficient `index(of:)`
  public func _customIndexOfEquatableElement(_ element: Iterator.Element) -> Index?? {
    let i = _partitionPoint(where: { !areInIncreasingOrder($0, element) })
    guard i != endIndex
      else { return .some(nil) }
    return !areInIncreasingOrder(element, self[i])
      ? .some(i)
      : .some(nil)
  }
  
  // Customization point for efficient `contains()`
  public func _customContainsEquatableElement(_ element: Iterator.Element) -> Bool? {
    return _customIndexOfEquatableElement(element)! != nil
  }
}

extension MutableSortedCollection {
  public mutating func insert<S: Sequence>(contentsOf newElements: S)
    where S.Iterator.Element == Iterator.Element
  {
    for element in newElements {
      insert(element)
    }
  }
  
  public mutating func insert<S: Sequence>(sortedContentsOf newElements: S)
    where S.Iterator.Element == Iterator.Element
  {
    var i = startIndex
    for element in newElements {
      i = _stableInsertionPoint(for: element, startingAt: i)
      _insert(element, atSortedIndex: i)
    }
  }
  
  @discardableResult
  public mutating func remove(at i: Index) -> Iterator.Element {
    let element = self[i]
    removeSubrange(i..<index(after: i))
    return element
  }
  
  public mutating func removeAll(keepingCapacity: Bool) {
    removeSubrange(startIndex..<endIndex)
  }
  
  public mutating func _insert(_ element: Iterator.Element, atSortedIndex i: Index) {
    insert(element)
  }
}

// MARK: - Sorted Types

/// An always-sorted random-access array.
public struct SortedArray<Element>: MutableSortedCollection, RandomAccessCollection {
  public let areInIncreasingOrder: (Element, Element) -> Bool
  private var _elements: [Element]
  
  // Collection requirements
  
  public var startIndex: Int { return 0 }
  public var endIndex: Int { return _elements.count }
  public subscript(i: Int) -> Element {
    return _elements[i]
  }
  public func index(after i: Int) -> Int {
    return i + 1
  }
  public func index(before i: Int) -> Int {
    return i - 1
  }
  public func index(_ i: Int, offsetBy n: Int) -> Int {
    return i + n
  }
  public func distance(from start: Int, to end: Int) -> Int {
    return end - start
  }
  
  /// Creates a new `SortedArray` by sorting the given sequence using the given
  /// predicate.
  public init<S: Sequence>(sorting elements: S,
    areInIncreasingOrder: @escaping (Element, Element) -> Bool)
    where S.Iterator.Element == Element
  {
    self._elements = elements.sorted(by: areInIncreasingOrder)
    self.areInIncreasingOrder = areInIncreasingOrder
  }
  
  /// Creates a new `SortedArray` using the given sequence and predicate.
  ///
  /// `elements` must already be in order before calling
  /// `init(sorted:areInIncreasingOrder:)`.
  public init<S: Sequence>(sorted elements: S,
    areInIncreasingOrder: @escaping (Element, Element) -> Bool)
    where S.Iterator.Element == Element,
      S.SubSequence : Sequence,
      S.SubSequence.Iterator.Element == Element
  {
    precondition(elements.isSorted(by: areInIncreasingOrder),
                 "unsorted elements passed to init(sorted:...)")
    self._elements = Array(elements)
    self.areInIncreasingOrder = areInIncreasingOrder
  }
  
  public mutating func _insert(_ element: Element, atSortedIndex i: Int) {
    precondition(i == 0 || !areInIncreasingOrder(element, _elements[i - 1]),
                 "insert(_:atSortedIndex:) with incorrect index")
    precondition(i >= _elements.count - 1 || areInIncreasingOrder(element, _elements[i + 1]),
                 "insert(_:atSortedIndex:) with incorrect index")
    _elements.insert(element, at: i)
  }
  
  @discardableResult
  public mutating func insert(_ element: Element) -> Int {
    let i = stableInsertionPoint(for: element)
    _elements.insert(element, at: i)
    return i
  }
  
  public mutating func removeSubrange(_ bounds: Range<Int>) {
    _elements.removeSubrange(bounds)
  }
}

extension SortedArray where Element : Comparable {
  /// Creates a new `SortedArray` by sorting the given sequence.
  ///
  /// The new array uses the element type's default `<` ordering.
  public init<S: Sequence>(sorting elements: S)
    where S.Iterator.Element == Element
  {
    self = .init(sorting: elements, areInIncreasingOrder: <)
  }
  
  /// Creates a new `SortedArray` using the given sequence.
  ///
  /// The new array uses the element type's default `<` ordering.
  /// `elements` must already be in ascending order before calling
  /// `init(sorted:)`.
  public init<S: Sequence>(sorted elements: S)
    where S.Iterator.Element == Element,
      S.SubSequence : Sequence,
      S.SubSequence.Iterator.Element == Element
  {
    self = .init(sorted: elements, areInIncreasingOrder: <)
  }
}

/// A wrapper around an already-sorted collection.
public struct SortedView<Base: Collection> : SortedCollection
  where Base.Iterator.Element == Base.SubSequence.Iterator.Element
{
  typealias Element = Base.Iterator.Element
  public typealias Index = Base.Index
  
  public let areInIncreasingOrder: (Element, Element) -> Bool
  let base: Base
  
  // Collection requirements
  
  public var startIndex: Index { return base.startIndex }
  public var endIndex: Index { return base.endIndex }
  public func index(after i: Index) -> Index {
    return base.index(after: i)
  }
  public subscript(i: Index) -> Element {
    return base[i]
  }
  
  /// Creates a new `SortedView` using the given sequence and predicate.
  ///
  /// `elements` must already be in order before calling
  /// `init(sorted:areInIncreasingOrder:)`.
  public init(sorted elements: Base,
    areInIncreasingOrder: @escaping (Element, Element) -> Bool)
  {
    precondition(elements.isSorted(by: areInIncreasingOrder),
                 "unsorted elements passed to init(sorted:...)")
    self.base = elements
    self.areInIncreasingOrder = areInIncreasingOrder
  }
}

extension SortedView where Base.Iterator.Element : Comparable {
  /// Creates a new `SortedView` using the given sequence.
  ///
  /// The new wrapper uses the element type's default `<` ordering.
  /// `elements` must already be in ascending order before calling
  /// `init(sorted:)`.
  public init(sorted elements: Base) {
    self = .init(sorted: elements, areInIncreasingOrder: <)
  }
}

// MARK: - Sequence and Collection Extensions

extension Sequence
  where SubSequence : Sequence,
  SubSequence.Iterator.Element == Iterator.Element
{
  /// Returns `true` if the sequence is in order according to the given closure.
  public func isSorted(
    by areInIncreasingOrder: (Iterator.Element, Iterator.Element) -> Bool) -> Bool
  {
    for (a, b) in zip(self, self.dropFirst()) {
      if areInIncreasingOrder(b, a) {
        return false
      }
    }
    return true
  }
}

extension Collection {
  /// Returns a partitioning index `p`, where `predicate` returns `false` for
  /// all the elements up to, but not including `p`, and `true` for all
  /// elements from `p` onward.
  ///
  /// The collection must already be partitioned according to `predicate`.
  internal func _partitionPoint(
    where predicate: (Iterator.Element) throws -> Bool
    ) rethrows -> Index
  {
    return try _partitionPoint(where: predicate, startingAt: startIndex)
  }
  
  // Workaround for missing generic features.
  internal func _partitionPoint(
    where predicate: (Iterator.Element) throws -> Bool,
    startingAt i: Index
    ) rethrows -> Index
  {
    var n = distance(from: i, to: endIndex)
    var r = i
    
    while n > 0 {
      let half = n / 2
      let mid = index(r, offsetBy: half)
      if try predicate(self[mid]) {
        n = half
      }
      else {
        r = index(after: mid)
        n -= half + 1
      }
    }
    return r
  }
}