Elegant way to get the first n elements of a SequenceType

Question

Is there an elegant way (think one-liner) to get the first n elements of a SequenceType in Swift?

I could of course write a for-loop that terminates after n elements but that's a little bulky.

Note also that the solution should be able to deal with infinite sequences.

Answer 1

Isn't it exactly what mySequence.prefix(numberOfElements) does?

Answer 2

In Swift 2, you can create this as an extension:

extension SequenceType {
    func take(n: Int) -> [Generator.Element] {
        var result: [Generator.Element] = []
        var g = self.generate()
        for _ in 1...n {
            if let next = g.next() {
                result.append(next)
            } else {
                break
            }
        }
        return result
    }
}

In Swift 1, it would have to written as a function:

func take<Seq: SequenceType>(n: Int, xs: Seq) -> [Seq.Generator.Element] {
    var result: [Seq.Generator.Element] = []
    var g = xs.generate()
    for _ in 1...n {
        if let next = g.next() {
            result.append(next)
        } else {
            break
        }
    }
    return result
}

Note that in either case, SequenceType does not specify what happens if you call generate() more than once. It could return the same values (as in an Array). It could return different values (as in an audio data stream). It could return nothing at all. So the caller of take() may need some special knowledge about its impact on the sequence.