Swift Challenge #2

gistfile1.swift

// https://gist.github.com/Eiam8821/07a7372667879172e9df

enum TakeWhileState {
    case Continuing(Int)
    case Done
}

struct TakeWhileGenerator<T>: GeneratorType {
    let sequence: TakeWhile<T>
    var state: TakeWhileState
    
    init(sequence: TakeWhile<T>) {
        self.sequence = sequence
        self.state = 0 < sequence.array.count ? .Continuing(0) : .Done
    }
    
    mutating func next() -> T? {
        switch state {
        case let .Continuing(index):
            let item = sequence.array[index]
            if sequence.condition(item) {
                state = index + 1 < sequence.array.count ? .Continuing(index + 1) : .Done
                return item
            } else {
                state = .Done
                return nil
            }
        case .Done:
            return nil
        }
    }
}

struct TakeWhile<T>: SequenceType {
    let array: [T]
    let condition: T -> Bool
    
    init(array: [T], condition: T -> Bool) {
        self.array = array
        self.condition = condition
    }
    
    func generate() -> TakeWhileGenerator<T> {
        return TakeWhileGenerator(sequence: self)
    }
}

extension Array {
    func takeWhile(condition: Element -> Bool) -> [Element] {
        return Array(TakeWhile(array: self, condition: condition))
    }
}

for item in (TakeWhile(array: [1, 2, 3, 4, 5]) { $0 < 3 }) {
    println(item)
}

Array itself is a Sequence so you can use its generator instead of dealing with its indexes and worrying about state. Once you do that, you can make TakeWhile even more generic if you take any sequence and use the same generator logic.