Use AnySequence and anyGenerator in combination

Question

Another question asked, essentially, how to implement a take function which would return the first n elements of a sequence. My answer was:

struct TakeFromSequenceSequence<S:SequenceType> : SequenceType {
    var limit : Int
    var sequence : S

    func generate() -> AnyGenerator<S.Generator.Element> {
        var generator = sequence.generate()
        var limit = self.limit

        return anyGenerator {
            guard limit > 0 else {
                return nil
            }

            limit = limit - 1

            return generator.next()
        }
    }
}

extension SequenceType {
    func take(count:Int) -> TakeFromSequenceSequence<Self> {
        return TakeFromSequenceSequence(limit: count, sequence: self)
    }
}

but it seems like I ought to be able to use AnySequence and anyGenerator to do it all inline in my take function:

extension SequenceType {
    func take(count:Int) -> AnySequence<Self.Generator.Element> {
        // cannot invoke initializer for type 'AnySequence<_>' with an argument list of type '(() -> _)'
        return AnySequence({
            var generator = self.generate()
            var limit = count

            // cannot invoke 'anyGenerator' with an argument list of type '(() -> _)'
            return anyGenerator({
                guard limit > 0 else {
                    return nil
                }

                limit = limit - 1

                return generator.next()
            })
        })
    }
}

Unfortunately, this yields multiple typing errors, mostly (I think) because type inference is failing.

Anybody have any suggestions on how to get this (using AnySequence and anyGenerator inline) to work?

Answer 1

(The answer is now based on Swift 2.2/Xcode 7.3. A solution for Swift 2.1 can be found in the edit history.)

The type of the closure passed to the AnySequence init method must be specified explicitly:

extension SequenceType {
    func take(count:Int) -> AnySequence<Generator.Element> {

        return AnySequence { () -> AnyGenerator<Generator.Element> in
            var generator = self.generate()
            var limit = count
            return AnyGenerator {  
                guard limit > 0 else {
                    return nil
                }
                limit = limit - 1
                return generator.next()
            }
        }
    }
}

Note that the (redundant) Self. in Self.Generator.Element is omitted, otherwise it does not compile.

Example:

let sequence = [1,2,3,4,5].take(2)
print(Array(sequence)) // [1, 2]
print(Array(sequence)) // [1, 2]

Alternatively, the method can be defined as

extension SequenceType {
    func take(count:Int) -> AnySequence<Generator.Element> {

        var generator = self.generate()
        var limit = count

        return AnySequence {
            return AnyGenerator {  
                guard limit > 0 else {
                    return nil
                }
                limit = limit - 1
                return generator.next()
            }
        }
    }
}

Now the closure passed to the AnySequence init method is a "single-expression closure" and the type is inferred by the compiler.

But – as David Berry noted – the created sequence then behaves differently, the generate() method cannot be called repeatedly with identical results:

let sequenceX = [1,2,3,4,5].take(2)
print(Array(sequenceX)) // [1, 2]
print(Array(sequenceX)) // []

This is permitted behavior, as stated in the SequenceType protocol reference:

... It is not correct to assume that a sequence will either be "consumable" and will resume iteration, or that a sequence is a collection and will restart iteration from the first element. A conforming sequence that is not a collection is allowed to produce an arbitrary sequence of elements from the second generator.

So one can choose among these implementations, dependent on the desired behavior.