Make defining a collection as easy as pie

pie.swift

//===----------------------------------------------------------------------===//
//===--- Defining Collections, easy as Pie --------------------------------===//
//===----------------------------------------------------------------------===//
struct Digits : PieCollection {
  let startState: Optional = 0
  func iterate(from state: Int) -> (nextState: Int?, element: Int) {
    return (nextState: state+1 < 10 ? state+1 : nil, element: state)
  }
}

func testDigits() {
  let d = Digits()
  var i = d.startIndex
  while i != d.endIndex {
    print(d[i])
    d.formIndex(after: &i)
  }
}

struct Fib : PieCollection {
  let startState: Optional = (0, 1)
  func iterate(from state: (Int, Int)) -> (nextState: (Int, Int)?, element: Int) {
    let next = state.0 + state.1
    return (nextState: state.1 < 1000 ? (state.1, next) : nil, element: state.0)
  }
}

func testFib() {
  let fib = Fib()
  for i in fib.indices {
    print(fib[i])
  }
  print("index of 21 in fib:", fib.index(of: 21)!)
  var i = fib.startIndex
  while i != fib.endIndex {
    fib.formIndex(after: &i)
  }
  print("index of 22 in fib:", fib.index(of: 22))
}

//===--- Implementation of PieCollection ----------------------------------===//

// Workaround for compiler limitation; you can ignore this
protocol _Pie {
  associatedtype IterationState
  associatedtype Element
  
  var startState : IterationState? { get }
  func iterate(from: IterationState) -> (nextState: IterationState?, element: Element)
}

/// Makes defining a Collection easy as Pie
///
/// just fulfill the requirements below
protocol PieCollection : _Pie, Collection {
  associatedtype IterationState
  associatedtype Element
  
  var startState : IterationState? { get }
  func iterate(from: IterationState) -> (nextState: IterationState?, element: Element)
}

struct PieIndex<C: _Pie>  : Comparable {
  init(state: C.IterationState?, offset: Int) {
    self.state = state
    self.offset = offset
  }
  
  var state: C.IterationState?
  var offset: Int
}

func < <T>(lhs: PieIndex<T>, rhs: PieIndex<T>) -> Bool {
  return lhs.state != nil && (rhs.state == nil || lhs.offset < rhs.offset)
}

func == <T>(lhs: PieIndex<T>, rhs: PieIndex<T>) -> Bool {
  return lhs.state == nil
  ? rhs.state == nil
  : rhs.state != nil && lhs.offset == rhs.offset
}

// Not really needed, but should be more efficient than the default
// you get from Collection.
struct PieIterator<C: _Pie> : IteratorProtocol {
  init(_ elements: C) {
    self.elements = elements
    self.state = elements.startState
  }
  
  mutating func next() -> C.Element? {
    if state == nil { return nil }
    let e: Element
    (state, e) = elements.iterate(from: state!)
    return e
  }
  
  var state: C.IterationState?
  let elements: C
}

extension PieCollection {
  var startIndex : PieIndex<Self> {
    return PieIndex(state: startState, offset: 0)
  }
  
  var endIndex : PieIndex<Self> {
    return PieIndex(state: nil, offset: 0)
  }

  func index(after i: PieIndex<Self>) -> PieIndex<Self> {
    return PieIndex(state: iterate(from: i.state!).nextState, offset: i.offset + 1)
  }

  subscript(i: PieIndex<Self>) -> Element {
    return iterate(from: i.state!).element
  }
  
  // Not really needed, but should be more efficient than the default
  // you get from Collection.
  func makeIterator() -> PieIterator<Self> {
    return PieIterator(self)
  }
}


//===--- Run the tests ----------------------------------------------------===//


testDigits()
print("--")
testFib()