Cycle.swift

/// The iterator for `LazyCycleSequence`.
public struct LazyCycleIterator <BaseSequence: Sequence>: IteratorProtocol {
  public typealias Element = BaseSequence.Iterator.Element

  private let baseSequence: BaseSequence
  private var baseIterator: BaseSequence.Iterator

  internal init (_ baseSequence: BaseSequence) {
    self.baseSequence = baseSequence
    self.baseIterator = self.baseSequence.makeIterator()
  }

  public mutating func next () -> Element? {
    var next = self.baseIterator.next()

    if next == nil {
      self.baseIterator = self.baseSequence.makeIterator()
      next = self.baseIterator.next()
    }

    return next
  }
}

/// A lazily evaluated `Sequence` that infinitely repeats the elements of the base `Sequence`.
///
/// - Attention: To create an instance of `LazyCycleSequence`, use the `cycle` property.
public struct LazyCycleSequence <BaseSequence: Sequence>: LazySequenceProtocol {
  public typealias Iterator = LazyCycleIterator<BaseSequence>
  public typealias Element = Iterator.Element

  internal let baseSequence: BaseSequence

  internal init (_ baseSequence: BaseSequence) {
    self.baseSequence = baseSequence
  }

  public func makeIterator () -> Iterator {
    return Iterator(self.baseSequence)
  }
}

public extension LazyCycleSequence {
  /// Creates a lazily evaluated Sequence that infinitely repeats the elements of the base Sequence.
  ///
  ///     print(Array([0, 1, 2].cycle.prefix(4)))
  ///     // [0, 1, 2, 0]
  var cycle: LazyCycleSequence {
    return LazyCycleSequence(self.baseSequence)
  }
}

public extension LazySequence {
  /// Creates a lazily evaluated Sequence that infinitely repeats the elements of the base Sequence.
  ///
  ///     print(Array([0, 1, 2].cycle.prefix(4)))
  ///     // [0, 1, 2, 0]
  var cycle: LazyCycleSequence<Elements> {
    return LazyCycleSequence(self.elements)
  }
}

public extension Sequence {
  /// Creates a lazily evaluated Sequence that infinitely repeats the elements of the base Sequence.
  ///
  ///     print(Array([0, 1, 2].cycle.prefix(4)))
  ///     // [0, 1, 2, 0]
  var cycle: LazyCycleSequence<Self> {
    return LazyCycleSequence(self)
  }
}

import XCTest
@testable import SequenceExtensions

/// - Note: Since `cycle` is a newly defined property without any overloads we only check the type and value for the
///         base `Sequence`.
class CycleTests: XCTestCase {
  let baseSequence = sequence(first: 0, next: {
    if $0 < 2 {
      return $0 + 1
    }

    return nil
  })

  typealias BaseSequence = UnfoldSequence<Int, (Int?, Bool)>

  func testShouldReturnCorrectType () {
    XCTAssertTrue(type(of: self.baseSequence.cycle) == LazyCycleSequence<BaseSequence>.self)
  }

  func testShouldReturnCorrectValue () {
    XCTAssertEqual(Array(self.baseSequence.cycle.prefix(0)), [])
    XCTAssertEqual(Array(self.baseSequence.cycle.prefix(1)), [0])
    XCTAssertEqual(Array(self.baseSequence.cycle.prefix(4)), [0, 1, 2, 0])
  }
}