import XCTest

extension Collection {

  public func cycledCustom() -> CycleSequence<Self> {
    return CycleSequence(self)
  }

  public func cycledRepeat() -> FlattenSequence<Repeated<Self>> {
    return repeatElement(self, count: self.isEmpty ? 0 : .max).joined()
  }

  public func cycledUnfoldFirst() -> FlattenSequence<UnfoldFirstSequence<Self>> {
    return sequence(first: self, next: { $0.isEmpty ? nil : $0 }).joined()
  }

  public func cycledUnfoldState() -> UnfoldSequence<Element, Iterator> {
    return sequence(state: makeIterator(), next: { iterator in
      if let nextElement = iterator.next() {
        return nextElement
      }
      iterator = self.makeIterator() // NOTE: captures `self`.
      return iterator.next()
    })
  }
}

public struct CycleSequence<C: Collection> : Sequence, IteratorProtocol {

  internal let _collection: C
  internal var _iterator: C.Iterator

  internal init(_ collection: C) {
    _collection = collection
    _iterator = collection.makeIterator()
  }

  public mutating func next() -> C.Element? {
    if let nextElement = _iterator.next() {
      return nextElement
    }
    _iterator = _collection.makeIterator();
    return _iterator.next()
  }
}

class CycleBenchmark: XCTestCase {

  let _prefix = 1_000_000
  let _source = 0 ... 999

  func testCustom() {
    self.measure {
      var count = 0
      for _ in _source.cycledCustom().prefix(_prefix) {
        count += 1
      }
      XCTAssertEqual(count, _prefix)
    }
  }

  func testRepeat() {
    self.measure {
      var count = 0
      for _ in _source.cycledRepeat().prefix(_prefix) {
        count += 1
      }
      XCTAssertEqual(count, _prefix)
    }
  }

  func testUnfoldFirst() {
    self.measure {
      var count = 0
      for _ in _source.cycledUnfoldFirst().prefix(_prefix) {
        count += 1
      }
      XCTAssertEqual(count, _prefix)
    }
  }

  func testUnfoldState() {
    self.measure {
      var count = 0
      for _ in _source.cycledUnfoldState().prefix(_prefix) {
        count += 1
      }
      XCTAssertEqual(count, _prefix)
    }
  }
}