Custom iterator to infinitely iterate collection in a loop mode

Question

I am looking for iterator to infinitely iterate collection in a loop mode. So that when end index of collection is reached, then iterator should return element at start index.

The following solution seems working, but I hope it can be made in a better way.

public struct LoopIterator<T: Collection>: IteratorProtocol {

   private let collection: T
   private var startIndexOffset: T.IndexDistance

   public init(collection: T) {
      self.collection = collection
      startIndexOffset = 0
   }

   public mutating func next() -> T.Iterator.Element? {
      guard !collection.isEmpty else {
         return nil
      }
      let index = collection.index(collection.startIndex, offsetBy: startIndexOffset)
      startIndexOffset += T.IndexDistance(1)
      if startIndexOffset >= collection.count {
         startIndexOffset = 0
      }
      return collection[index]
   }
}

extension Array {
   func makeLoopIterator() -> LoopIterator<Array> {
      return LoopIterator(collection: self)
   }
}

// Testing...
// Will print: 1, 2, 3, 1, 2, 3
var it = [1, 2, 3].makeLoopIterator()
for _ in 0..<6 {
   print(it.next())
}

Is it a right way to do custom iterator? What can be improved?

Thanks!

Answer 1

In Swift 3 (which you're using), indexes are intended to be advanced by the collection itself. With that, you can simplify this as follows:

public struct LoopIterator<Base: Collection>: IteratorProtocol {

    private let collection: Base
    private var index: Base.Index

    public init(collection: Base) {
        self.collection = collection
        self.index = collection.startIndex
    }

    public mutating func next() -> Base.Iterator.Element? {
        guard !collection.isEmpty else {
            return nil
        }

        let result = collection[index]
        collection.formIndex(after: &index) // (*) See discussion below 
        if index == collection.endIndex {
            index = collection.startIndex
        }
        return result
    }
}

Now we simply move the index forward, and if it now points to the end, reset it to the beginning. No need for count or IndexDistance.

Note that I've used formIndex here, which exists to improve performance in somewhat obscure cases (specifically around AnyIndex) since your Iterator works on any Collection (and therefore any Index). The simpler version would be index = collection.index(after: index), and that may be better in most cases.

For all the gory details on Swift 3 indices, see SE-0065.