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!

For future reference, questions about improvements to working code are probably better suited for the **Code Review** stack exchange. http://codereview.stackexchange.com/ — Dan Beaulieu, Jul 16 '16 at 16:36
IMO, codereview.se is best suited for review of code that itself would be of little general interest and unlikely to be searched for explicitly (i.e. the best practices are of interest, not the code). This question is about a generally interesting piece of code, touches specifically on a topic of broad interest (correct use of Swift 3 indices), and is likely to be searched for in the future. Had the poster not included code, they certainly would have received "What have you tried?" comments. If they do include code, it seems unfair to redirect them to codereview. — Rob Napier, Jul 16 '16 at 16:55

Rob Napier · Accepted Answer · 2016-07-16T16:47:04.900

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.

Imanou Petit · Answer 2 · 2020-05-01T17:26:33.300

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]

Custom iterator to infinitely iterate collection in a loop mode

2 Answers2

#1. Using AnyIterator

#2. Using AnySequence

#1. Using `AnyIterator`

#2. Using `AnySequence`