coroutines.swift

import Cocoa

enum CoroutineState {
    case Fresh, Running, Blocked, Canceled, Done
}

struct CoroutineCancellation: ErrorType {}

class CoroutineImpl<InputType, YieldType> {
    let body: (yield: YieldType throws -> InputType) throws -> Void
    
    var state = CoroutineState.Fresh
    
    let lock = NSCondition()
    
    var inputValue: InputType?
    var yieldValue: YieldType?
    
    init(body: (yield: YieldType throws -> InputType) throws -> Void) {
        self.body = body
    }
    
    func getNext(value: InputType) -> YieldType? {
        lock.lock()
        inputValue = value
        unblock()
        while state == .Running {
            lock.wait()
        }
        let result = yieldValue
        yieldValue = nil
        
        lock.unlock()
        
        return result
    }
    
    func cancel() {
        lock.lock()
        state = .Canceled
        lock.signal()
        lock.unlock()
    }
    
    private func unblock() {
        switch state {
        case .Fresh:
            state = .Running
            dispatch_async(dispatch_get_global_queue(0, 0), {
                do {
                    try self.body(yield: self.yield)
                } catch {}
                self.lock.lock()
                self.state = .Done
                self.lock.signal()
                self.lock.unlock()
            })
        case .Running:
            preconditionFailure("Must never call unblock() while running")
        case .Blocked:
            state = .Running
            lock.signal()
        case .Canceled:
            break
        case .Done:
            break
        }
    }
    
    private func yield(value: YieldType) throws -> InputType {
        lock.lock()
        self.yieldValue = value
        state = .Blocked
        lock.signal()
        while state == .Blocked {
            lock.wait()
        }
        let canceled = state == .Canceled
        let input = inputValue
        inputValue = nil
        lock.unlock()
        
        if canceled {
            throw CoroutineCancellation()
        }
        return input!
    }
}

class Coroutine<InputType, YieldType> {
    let impl: CoroutineImpl<InputType, YieldType>
    
    init(_ body: (yield: YieldType throws -> InputType) throws -> Void) {
        impl = CoroutineImpl(body: body)
    }
    
    deinit {
        impl.cancel()
    }
    
    func getNext(value: InputType) -> YieldType? {
        return impl.getNext(value)
    }
}

func coroutine<InputType, YieldType>(type: (InputType, YieldType).Type, _ body: (yield: YieldType throws -> InputType) throws -> Void) -> (InputType -> YieldType?){
    let obj = Coroutine(body)
    return obj.getNext
}

var c: Coroutine<Void, Int>? = Coroutine<Void, Int>({ yield in
    for i in 0..<10 {
        try yield(i)
    }
})

c?.getNext()
c?.getNext()
c?.getNext()
c?.getNext()
c = nil

var d: (Void -> Int?)? = coroutine((Void, Int).self, { yield in
    for i in 20..<30 {
        try yield(i)
    }
})

d?()
d?()
d?()
d = nil

let summer = coroutine((Int, Int).self, { yield in
    var sum = 0
    while true {
        let value = try yield(sum)
        print(value)
        sum += value
    }
})

summer(0)
summer(1)
summer(2)
summer(3)
summer(4)
summer(5)

sleep(1)