//! The state monad transformer.

let { Alternative, or, empty } = import! std.alternative
let { Applicative, wrap, (<*>) } = import! std.applicative
let { (>>), (<<), ? } = import! std.function
let { Functor, map } = import! std.functor
let { Monad, (>>=) } = import! std.monad
let { Transformer } = import! std.transformer

type StateOut s a = { value : a, state : s }
type WrStateOut s m a = m { value : a, state : s }

type StateT s m a = s -> m { value : a, state : s }

let map_sout f st : (a -> b) -> StateOut s a -> StateOut s b =
    { value = f st.value, state = st.state }

let functor : [Functor m] -> Functor (StateT s m) =
    let stmap f sr : (a -> b) -> StateT s m a -> StateT s m b =
        map (map_sout f) << sr

    { map = stmap }

// the typechecker can't find map and Functor m without help
let applicative ?mo : [Monad m] -> Applicative (StateT s m) =
    let apply srf sr : StateT s m (a -> b) -> StateT s m a -> StateT s m b = \state ->
        srf state
            >>= \fout ->
                let { value = f, state = state' } = fout
                mo.applicative.functor.map (map_sout f) (sr state')

    let stwrap value : a -> StateT s m a = \state ->
        wrap { value, state }

    { functor = functor ?mo.applicative.functor, apply, wrap = stwrap }

let monad : [Monad m] -> Monad (StateT s m) =
    let flat_map f sr : (a -> StateT s m b) -> StateT s m a -> StateT s m b = \state ->
        sr state
            >>= \sout ->
                let { value, state = state' } = sout
                f value state'

    { applicative, flat_map }

let transformer : Transformer (StateT s) =
    let wrap_monad ma : [Monad m] -> m a -> StateT s m a = \state ->
        ma >>= \value -> wrap { value, state }

    { /* monad, */ wrap_monad }

let alternative : [Monad m] -> [Alternative m] -> Alternative (StateT s m) =
    let stempty = transformer.wrap_monad empty
    let stor sra srb = or << sra <*> srb

    { applicative, empty = stempty, or = stor }

let put value : [Monad m] -> s -> StateT s m () = \state ->
    wrap { value = (), state = value }

let get : [Monad m] -> StateT s m s = \state ->
    wrap { value = state, state }

let gets f : [Monad m] -> (s -> a) -> StateT s m a =
    get >>= (wrap << f)

let modify f : [Monad m] -> (s -> s) -> StateT s m () =
    get >>= (put << f)

let run_state_t f state : StateT s m a -> s -> m { value : a, state : s } =
    f state

let eval_state_t f state : [Functor m] -> StateT s m a -> s -> m a =
    map (\x -> x.value) (run_state_t f state)

let exec_state_t f state : [Functor m] -> StateT s m a -> s -> m s =
    map (\x -> x.state) (run_state_t f state)

{
    StateT,

    applicative,
    functor,
    monad,
    transformer,
    alternative,

    put,
    get,
    gets,
    modify,
    run_state_t,
    eval_state_t,
    exec_state_t,
}