Stateモナド @Swift

@kazu_yamamotoさんのStateモナド解説記事をSwiftで書くとこんな感じというメモ。

indirect enum List<A> {
    case Nil
    case Cons(A, List<A>)
}

typealias Counter = Int

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

    static func pure(x: A) -> State {
        return State(runState: { (x, $0) })
    }

    func bind<B>(g: A -> State<S, B>) -> State<S, B> {
        return State<S, B>(runState: {
            let (a, s) = self.runState($0)
            let s2 = g(a)
            return s2.runState(s)
        })
    }
}

infix operator >>- {
    associativity left
    precedence 100
}

func >>- <S, A, B>(f: State<S, A>, g: A -> State<S, B>) -> State<S, B> {
    return f.bind(g)
}

func cons0<A>(x: A) -> List<A> -> State<Counter, List<A>> {
    return { xs in
        return State(runState: { (.Cons(x, xs), $0 + 1) })
    }
}

func run<A>(state: State<Counter, A>) -> Counter -> (A, Counter) {
    return { state.runState($0) }
}

let incr1: State<Int, Int> = State(runState: { (999, $0 + 1) })
func cons1<A>(x: A) -> List<A> -> State<Counter, List<A>> {
    return { xs in
        return incr1 >>- { _ in .pure(.Cons(x, xs)) }
    }
}

let incr2: State<Counter, Counter> = State(runState: { ($0, $0 + 1) })
func cons2<A>(x: A) -> List<(Counter, A)> -> State<Counter, List<(Counter, A)>> {
    return { xs in
        return incr2 >>- { .pure( .Cons(($0, x), xs) ) }
    }
}

func get() -> State<Counter, Counter> {
    return State(runState: { ($0, $0) })
}

func put(x: Counter) -> State<Counter, Counter> {
    return State(runState: { ($0, x) })
}

func cons3<A>(x: A) -> List<(Counter, A)> -> State<Counter, List<(Counter, A)>> {
    return { xs in
        return get()
            >>- { put($0 + 1) }
            >>- { .pure(.Cons(($0, x), xs)) }
    }
}

run(cons3("a")(.Nil) >>- cons3("b") >>- cons3("c"))(0)

何かに使えないものだろうか。おとなしく変数に入れろって話かもしれませんが。。。