Monad composition (Cont · State)

Question

I'm studying monad composition. While I already understand how to compose, say, Async and Result as performed here I'm struggling in composing the Continuation Monad and the State Monad.

Starting from a basic State Monad implementation and aState-based-Stack for testing purposes:

type State<'State,'Value> = State of ('State -> 'Value * 'State)

module State =
    let runS (State f) state = f state

    let returnS x =
        let run state =
            x, state
        State run

    let bindS f xS =
        let run state =
            let x, newState = runS xS state
            runS (f x) newState
        State run

    let getS =
        let run state = state, state
        State run

    let putS newState =
        let run _ = (), newState
        State run

    type StateBuilder()=
        member __.Return(x) = returnS x
        member __.Bind(xS,f) = bindS f xS

    let state = new StateBuilder()

module Stack =
    open State

    type Stack<'a> = Stack of 'a list

    let popStack (Stack contents) = 
        match contents with
        | [] -> failwith "Stack underflow"
        | head::tail ->     
            head, (Stack tail)

    let pushStack newTop (Stack contents) = 
        Stack (newTop::contents)

    let emptyStack = Stack []

    let getValue stackM = 
        runS stackM emptyStack |> fst

    let pop() = state {
        let! stack = getS
        let top, remainingStack = popStack stack
        do! putS remainingStack 
        return top }

    let push newTop = state {
        let! stack = getS
        let newStack = pushStack newTop stack
        do! putS newStack 
        return () }

Then having also a basic implementation of a Continuation Monad :

type Cont<'T,'r> = (('T -> 'r) -> 'r)

module Continuation =
    let returnCont x = (fun k -> k x)
    let bindCont f m = (fun k -> m (fun a -> f a k))
    let delayCont f = (fun k -> f () k)
    let runCont (c:Cont<_,_>) cont = c cont
    let callcc (f: ('T -> Cont<'b,'r>) -> Cont<'T,'r>) : Cont<'T,'r> =
        fun cont -> runCont (f (fun a -> (fun _ -> cont a))) cont

    type ContinuationBuilder() =
        member __.Return(x) = returnCont x
        member __.ReturnFrom(x) = x
        member __.Bind(m,f) = bindCont f m
        member __.Delay(f) = delayCont f
        member this.Zero () = this.Return ()

    let cont = new ContinuationBuilder()

I'm trying to compose it like this :

module StateK =
    open Continuation

    let runSK (State f) state = cont { return f state }
    let returnSK x = x |> State.returnS |> returnCont

    let bindSK f xSK = cont {
        let! xS = xSK
        return (State.bindS f xS) }

    let getSK k =
        let run state = state, state
        State run |> k

    let putSK newState = cont {
        let run _ = (), newState
        return State run }

    type StateContinuationBuilder() =
        member __.Return(x) = returnSK x
        member __.ReturnFrom(x) = x
        member __.Bind(m,f) = bindSK f m
        member this.Zero () = this.Return () 

    let stateK = new StateContinuationBuilder()

While this compiles and seems right (as far as a mechanically-following-steps-composition goes) I'm not able to implement a StateK-based-Stack. So far I have this, but it is totally wrong:

module StackCont =
    open StateK

    type Stack<'a> = Stack of 'a list

    let popStack (Stack contents) =  stateK {
        match contents with
        | [] -> return failwith "Stack underflow"
        | head::tail ->     
            return head, (Stack tail) }

    let pushStack newTop (Stack contents) = stateK {
        return Stack (newTop::contents) }

    let emptyStack = Stack []

    let getValue stackM = stateK {
        return runSK stackM emptyStack |> fst }

    let pop() = stateK {
        let! stack = getSK
        let! top, remainingStack = popStack stack
        do! putSK remainingStack 
        return top }

    let push newTop = stateK {
        let! stack = getSK
        let! newStack = pushStack newTop stack
        do! putSK newStack 
        return () }

Some help to understand why and how is more than welcome. If there is some reading material you can point to, it will also work.

********* EDIT after AMieres comment **************

New bindSK implementation trying to keep signatures right.

type StateK<'State,'Value,'r> = Cont<State<'State,'Value>,'r>

module StateK =

    let returnSK x :  StateK<'s,'a,'r> = x |> State.returnS |> Continuation.returnCont
    let bindSK (f : 'a ->  StateK<'s,'b,'r>) 
        (m : StateK<'s,'a,'r>) :  StateK<'s,'b,'r> =
        (fun cont ->
            m (fun (State xS) ->
                let run state =
                    let x, newState = xS state
                    (f x) (fun (State k) -> k newState)
                cont (State run)))

Nevertheless, the type 'r has been constrained to be 'b * 's I have tried to remove the constraint but I haven't yet been able to do it

Answer 1

I have not been able to solve it either.

I can only give you a tip that may help you understand it better. Replace generic types for regular types, for instance instead of:

let bindSK (f : 'a ->  StateK<'s,'b,'r>) 
    (m : StateK<'s,'a,'r>) :  StateK<'s,'b,'r> =
    (fun cont ->
        m (fun (State xS) ->
            let run state =
                let x, newState = xS state
                (f x) (fun (State k) -> k newState)
            cont (State run)))

replace 's with string, 'a with int, 'b with char and 'r with float

let bindSK (f : int ->  StateK<string,char,float>) 
    (m : StateK<string,int,float>) :  StateK<string,char,float> =
    (fun cont ->
        m (fun (State xS) ->
            let run state =
                let x, newState = xS state
                (f x) (fun (State k) -> k newState)
            cont (State run)))

that way is easier to see that

• k is string -> char * string
• so k newState is char * string
• (f x) is (State<string,char> -> float) -> float
• and m is (State<string,int> -> float) -> float

so they are not compatible.

Answer 2

I read more and it comes out that the correct type for a "ContinuousState" is 's -> Cont<'a * 's, 'r>

So I re-implemented the StateK monad with this signatures and all flew naturally.

Here is the code (I added mapSK and applySK for completeness):

type Cont<'T,'r> = (('T -> 'r) -> 'r)

let returnCont x = (fun k -> k x)
let bindCont f m = (fun k -> m (fun a -> f a k))
let delayCont f = (fun k -> f () k)

type ContinuationBuilder() =
    member __.Return(x) = returnCont x
    member __.ReturnFrom(x) = x
    member __.Bind(m,f) = bindCont f m
    member __.Delay(f) = delayCont f
    member this.Zero () = this.Return ()

let cont = new ContinuationBuilder()

type StateK<'State,'Value,'r> = StateK of ('State -> Cont<'Value * 'State, 'r>)

module StateK =
    let returnSK x =
        let run state = cont {
            return x, state
        }
        StateK run

    let runSK (StateK fSK : StateK<'s,'a,'r>) (state : 's) : Cont<'a * 's, _> = cont {
        return! fSK state }

    let mapSK (f : 'a -> 'b) (m : StateK<'s,'a,'r>) : StateK<'s,'b,'r> =
            let run state = cont {
                let! x, newState = runSK m state
                return f x, newState  }
            StateK run

    let bindSK (f : 'a -> StateK<'s,'b,'r>) (xSK : StateK<'s,'a,'r>) : (StateK<'s,'b,'r>) =
        let run state = cont {
            let! x, newState = runSK xSK state
            return! runSK (f x) newState }
        StateK run

    let applySK (fS : StateK<'s, 'a -> 'b, 'r>) (xSK : StateK<'s,'a,'r>) : StateK<'s,'b,'r> =
        let run state = cont {
            let! f, s1 = runSK fS state
            let! x, s2 = runSK xSK s1
            return f x, s2 }
        StateK run        

    let getSK =
        let run state = cont { return state, state }
        StateK run

    let putSK newState =
        let run _ = cont { return (), newState }
        StateK run

    type StateKBuilder() =
        member __.Return(x) = returnSK x
        member __.ReturnFrom (x) = x
        member __.Bind(xS,f) = bindSK f xS
        member this.Zero() = this.Return ()

    let stateK = new StateKBuilder()

module StackCont =
    open StateK

    type Stack<'a> = Stack of 'a list

    let popStack (Stack contents) = 
        match contents with
        | [] -> failwith "Stack underflow"
        | head::tail ->     
            head, (Stack tail)

    let pushStack newTop (Stack contents) = 
        Stack (newTop::contents)

    let emptyStack = Stack []

    let getValueSK stackM = cont {
        let! f = runSK stackM emptyStack 
        return f |> fst }

    let pop() = stateK {
        let! stack = getSK
        let top, remainingStack = popStack stack
        do! putSK remainingStack 
        return top }

    let push newTop = stateK {
        let! stack = getSK
        let newStack = pushStack newTop stack
        do! putSK newStack 
        return () }

open StateK
open StackCont

let helloWorldSK = (fun () -> stateK {
    do! push "world"
    do! push "hello"
    let! top1 = pop()
    let! top2 = pop()
    let combined = top1 + " " + top2 
    return combined
})

let helloWorld = getValueSK (helloWorldSK ()) id
printfn "%s" helloWorld