R-Style Boolean Sequence Selectors in Swift

boolean-sequences.swift

extension Sequence where Element: Comparable {
  static func ==(_ lhs: Self, _ rhs: Element) -> LazyMapSequence<Self, Bool> {
    lhs.lazy.map { $0 == rhs }
  }

  static func !=(_ lhs: Self, _ rhs: Element) -> LazyMapSequence<Self, Bool> {
    lhs.lazy.map { $0 != rhs }
  }

  static func <(_ lhs: Self, _ rhs: Element) -> LazyMapSequence<Self, Bool> {
    lhs.lazy.map { $0 < rhs }
  }

  static func <=(_ lhs: Self, _ rhs: Element) -> LazyMapSequence<Self, Bool> {
    lhs.lazy.map { $0 <= rhs }
  }

  static func >(_ lhs: Self, _ rhs: Element) -> LazyMapSequence<Self, Bool> {
    lhs.lazy.map { $0 > rhs }
  }

  static func >=(_ lhs: Self, _ rhs: Element) -> LazyMapSequence<Self, Bool> {
    lhs.lazy.map { $0 >= rhs }
  }
}

struct Broadcasted<Base: Collection>: Collection where Base.Index: BinaryInteger {

  typealias Element = Base.Element
  typealias Index = Base.Index

  private var base: Base
  private var count_: Index

  var startIndex: Index { 0 }
  var endIndex: Index { count_ }

  func index(after i: Index) -> Index { i + 1 }

  subscript(index: Index) -> Element {
    base[index % Index(base.count)]
  }

  init(_ base: Base, count: Index) {
    self.base = base
    self.count_ = count
  }
}

extension Collection {
  subscript<C>(selection: C) -> [Element] where C: Collection, C.Index: BinaryInteger, C.Element == Bool {
    zip(Broadcasted(selection, count: C.Index(count)), self)
      .filter(\.0).map(\.1)
  }
}