Cross product collection adapter.

CrossProduct.swift

struct CrossProduct2<Base0: Collection, Base1: Collection>: Collection {
  public private(set) var base0: Base0
  public private(set) var base1: Base1

  init(_ base0: Base0, _ base1: Base1) {
    self.base0 = base0
    self.base1 = base1
  }

  struct Index: Comparable {
    public fileprivate(set) var base0: Base0.Index
    public fileprivate(set) var base1: Base1.Index
    public static func < (a: Self, b: Self) -> Bool { (a.base0, a.base1) < (b.base0, b.base1) }
  }

  var startIndex: Index {
    .init(base0: base0.startIndex, base1: base1.startIndex)
  }

  var endIndex: Index {
    base0.isEmpty || base1.isEmpty ? startIndex : .init(base0: base0.endIndex, base1: base1.endIndex)
  }

  func formIndex(after i: inout Index) {
    base0.formIndex(after: &i.base0)
    if i.base0 != base0.endIndex { return }
    base1.formIndex(after: &i.base1)
    if i.base1 != base1.endIndex { i.base0 = base0.startIndex }
  }

  func index(after i: Index) -> Index {
    var j = i
    formIndex(after: &j)
    return j
  }

  subscript(i: Index) -> (Base0.Element, Base1.Element) {
    (base0[i.base0], base1[i.base1])
  }
}

extension CrossProduct2: BidirectionalCollection
  where Base0: BidirectionalCollection, Base1: BidirectionalCollection
{
  func formIndex(before i: inout Index) {
    if i.base0 == base0.startIndex || i.base1 == base1.endIndex {
      base1.formIndex(before: &i.base1)
      i.base0 = base0.endIndex
    }
    base0.formIndex(before: &i.base0)
  }

  func index(before i: Index) -> Index {
    var j = i
    formIndex(before: &j)
    return j
  }
}

extension CrossProduct2: RandomAccessCollection
  where Base0: RandomAccessCollection, Base1: RandomAccessCollection
{
  private func offset(of i: Index) -> Int {
    base0.distance(from: base0.startIndex, to: i.base0)
      + base1.distance(from: base1.startIndex, to: i.base1) * base0.count
  }

  func index(_ i: Index, offsetBy n: Int) -> Index {
    if n == 0 { return i } // Handle the case where one base collection is empty
    let d = offset(of: i) + n

    return .init(
      base0: base0.index(base0.startIndex, offsetBy: d % base0.count),
      base1: base1.index(base1.startIndex, offsetBy: d / base0.count))
  }

  func distance(from a: Index, to b: Index) -> Int {
    return offset(of: b) - offset(of: a)
  }
}

print(Array(CrossProduct2(0..<4, EmptyCollection<Int>())))
print(Array(CrossProduct2(0..<4, (6..<9).map(String.init))))

print(Array(CrossProduct2(0..<4, EmptyCollection<Int>().reversed())))
print(Array(Array(CrossProduct2(0..<4, (6..<9).map(String.init)).reversed()).reversed()))