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.

Answer 2

With Swift 5, you can use one of the following examples in order to solve your problem.

---

#1. Using AnyIterator

As an alternative to creating a new type that conforms to IteratorProtocol, you can use AnyIterator. The following code, based on Rob Napier's answer, shows how to use it:

extension Array {

    func makeInfiniteLoopIterator() -> AnyIterator<Element> {
        var index = self.startIndex

        return AnyIterator({
            if self.isEmpty {
                return nil
            }

            let result = self[index]

            index = self.index(after: index)
            if index == self.endIndex {
                index = self.startIndex
            }

            return result
        })
    }

}

Usage:

let infiniteLoopIterator = [1, 2, 3].makeInfiniteLoopIterator()
for val in infiniteLoopIterator.prefix(5) {
    print(val)
}

/*
 prints:
 1
 2
 3
 1
 2
 */

let infiniteLoopIterator = [1, 2, 3].makeInfiniteLoopIterator()
let array = Array(infiniteLoopIterator.prefix(7))
print(array) // prints: [1, 2, 3, 1, 2, 3, 1]

let infiniteLoopIterator = [1, 2, 3].makeInfiniteLoopIterator()

let val1 = infiniteLoopIterator.next()
let val2 = infiniteLoopIterator.next()
let val3 = infiniteLoopIterator.next()
let val4 = infiniteLoopIterator.next()

print(String(describing: val1)) // prints: Optional(1)
print(String(describing: val2)) // prints: Optional(2)
print(String(describing: val3)) // prints: Optional(3)
print(String(describing: val4)) // prints: Optional(1)

---

#2. Using AnySequence

A similar approach is to use AnySequence:

extension Array {

    func makeInfiniteSequence() -> AnySequence<Element> {
        return AnySequence({ () -> AnyIterator<Element> in
            var index = self.startIndex

            return AnyIterator({
                if self.isEmpty {
                    return nil
                }

                let result = self[index]

                self.formIndex(after: &index) // alternative to: index = self.index(after: index)
                if index == self.endIndex {
                    index = self.startIndex
                }

                return result
            })
        })
    }

}

Usage:

let infiniteSequence = [1, 2, 3].makeInfiniteSequence()
for val in infiniteSequence.prefix(5) {
    print(val)
}

/*
 prints:
 1
 2
 3
 1
 2
 */

let infiniteSequence = [1, 2, 3].makeInfiniteSequence()
let array = Array(infiniteSequence.prefix(7))
print(array) // prints: [1, 2, 3, 1, 2, 3, 1]