statemonad.swift

statemonad.swift

// a simple Swift translation of the Haskell State monad
// taken from http://brandon.si/code/the-state-monad-a-tutorial-for-the-confused/


struct State<S,A> {
    let run: ((S)->(A,S))
}

let fromSToAAndS = { (v: Int) -> (String, Int) in
    if v%5 == 0 { return ("foo", v+1) }
    else { return ("bar", v+1) }
}

let stateIntStr = State(run: fromSToAAndS)

stateIntStr.run(5)

extension State {
    // return a = State $ \s -> (a, s)
    func rtn(_ v: A) -> State { return State { (v, $0)} }
    // and as a Swift constructor
    init(_ v: A) {
        self.run = { (v, $0) }
    }
    
    func flatMap<B>(_ fn: @escaping ((A) -> State<S, B>)) -> State<S,B> {
        return State<S,B> { (s:S) -> (B,S) in
            let (a, s2) = self.run(s)
            return fn(a).run(s2)
        }
    }
}

(stateIntStr.flatMap { (str) -> State<Int, String> in
    State { (s: Int) in (str+"'", s+1) }
}).run(5)