Take the first N elements from an infinite iterator or generator in Swift

Question

I use generators a lot in Python. Now I try to do something similar in Swift.

From an infinite factorial generator in Python.

def gen_factorial():
  current = 1
  N = 1
  while True:
    yield current
    N += 1
    current *= N

A Swift version would be

struct gen_factorial: Sequence, IteratorProtocol 
{
  var current = 1
  var N = 1
  mutating func next()-> Int?{
    defer {
      N += 1
      current *= N
    }
    return current
  }
}

I test it by taking the first 4 elements with

zip(1...4, 
    gen_factorial()
   ).map{$1}

And get 1, 2, 6, 24 as expected.

But when I want to write a helper function take(n, gen) to simplify it, I can't.

func take(_ n: Int, _ it: AnyIterator<Int>) -> [Int]{
  return zip(1...n, it).map {$1}
}

take(4, gen_factorial())

The error message is

error: cannot convert value of type 'gen_factorial' to expected argument type 'AnyIterator<Int>'

What type should it be, if not AnyIterator?

I am still new to Swift. Please help.

Answer 1

AnyIterator is a concrete type conforming to the Iterator (and Sequence ) protocol. In order to pass your iterator to that function you would have to wrap it into an AnyIterator :

print(take(4, AnyIterator(gen_factorial())))
// [1, 2, 6, 24]

The better solution would be to make the function generic so that it takes an arbitrary sequence type as argument:

func take<S: Sequence>(_ n: Int, _ it: S) -> [S.Element]{
    return zip(1...n, it).map {$1}
}

print(take(4, gen_factorial()))
// [1, 2, 6, 24]

Remarks:

• Instead of your helper function you can use the existing prefix(_ maxLength:) method of Sequence :

   print(Array(gen_factorial().prefix(4))) // [1, 2, 6, 24]

• The Swift naming conventions for types a upper camel-case.