OptionSet+Collection.swift

/// This is a useful index that can store a comparable element or the end of
/// a collection. Similar to https://github.com/apple/swift/pull/15193
enum IndexWithEnd<T : Comparable> : Comparable {
  case element(T)
  case end
  
  static func < (lhs: IndexWithEnd, rhs: IndexWithEnd) -> Bool {
    switch (lhs, rhs) {
    case (.element(let l), .element(let r)):
      return l < r
    case (.element, .end):
      return true
    case (.end, .element), (.end, .end):
      return false
    }
  }
}

/// This extension provides all the Collection requirements to an OptionSet
/// that specifies that its Index is the type above.
extension Collection
  where Self : OptionSet, Self.RawValue : FixedWidthInteger,
  Index == IndexWithEnd<Self.RawValue>
{
  func _rawBit(after value: RawValue) -> RawValue? {
    let shift = value.trailingZeroBitCount + 1
    let shiftedRawValue = rawValue >> shift
    if shiftedRawValue == 0 {
      return nil
    } else {
      return (1 as RawValue) << (shiftedRawValue.trailingZeroBitCount + shift)
    }
  }
  
  var startIndex: Index {
    return rawValue == 0
      ? .end
      : .element(1 << rawValue.trailingZeroBitCount)
  }
  
  var endIndex: Index {
    return .end
  }
  
  var isEmpty: Bool {
    return rawValue == 0
  }
  
  var count: Int {
    return rawValue.nonzeroBitCount
  }
  
  subscript(i: Index) -> Self {
    switch i {
    case .element(let e):
      return Self(rawValue: e)
    case .end:
      fatalError("Can't subscript with endIndex")
    }
  }
  
  func index(after i: Index) -> Index {
    switch i {
    case .element(let e):
      return _rawBit(after: e).map({ .element($0) }) ?? .end
    case .end:
      fatalError("Can't advance past endIndex")
    }
  }
}

/// Here's a minimal OptionSet - the only additions to the typical elements
/// are the Index type alias and the Collection conformance.
struct O : OptionSet, Collection {
  typealias Index = IndexWithEnd<Int>
  var rawValue: Int
  
  static var one = O(rawValue: 1)
  static var two = O(rawValue: 2)
  static var three = O(rawValue: 4)
  static var four = O(rawValue: 8)
  static var oneTwoThree: O = [O.one, O.two, O.three]
}

/// Now you can use collection APIs with any O instance.
let a: O = [.one, .four, .three]
a.first     // Optional(O(rawValue: 1))
Array(a)    // [O(rawValue: 1), O(rawValue: 4), O(rawValue: 8)]