Swift Challenge #2

swift-challenge-2.swift

struct TakeWhileGenerator<T: GeneratorType>: GeneratorType {
    typealias Element = T.Element

    var generator: T
    let test: (element: T.Element) -> Bool

    mutating func next() -> Element? {
        if let next = generator.next() {
            if test(element: next) {
                return next
            }
        }
        return nil
    }
}

struct TakeWhile<T : SequenceType>: SequenceType {
    typealias Generator = TakeWhileGenerator<T.Generator>

    let seq: T
    let test: (element: T.Generator.Element) -> Bool

    init(_ seq: T, test: (element: T.Generator.Element) -> Bool) {
        self.seq = seq
        self.test = test
    }

    func generate() -> Generator {
        return Generator(generator: self.seq.generate(), test: self.test)
    }
}

let array = [1,2,3,2,1]
Array(TakeWhile(array, {$0 < 3}))