State monads: Transitioning from one state type to another

Question

Let's say we have a stack of monads with a state monad transformer as the outer most transformer like this one:

-- | SEWT: Composition of State . Except . Writer monad transformers in that
-- order where Writer is the innermost transformer.
-- the form of the computation is: s -> (Either e (a, s), w)
newtype SEWT s e w m a = SEWT {
    _runSEWT :: StateT s (ExceptT e (WriterT w m)) a }
    deriving (Functor, Applicative, Monad,
              MonadState s, MonadError e, MonadWriter w)

-- | 'runSEWT': runs a 'SEWT' computation given an initial state.
runSEWT :: SEWT s e w m a -> s -> m (Either e (a, s), w)
runSEWT ev e = runWriterT $ runExceptT $ runStateT (_runSEWT ev) e

We then want to do, in some form: SEWT s e w m a -> s -> SEWT t e w m a. This is of course not possible using (>>=) or a do block since a state monad with s as state is not the same monad as one with t.

I can then conjure up something like this:

-- | 'sewtTransition': transitions between one 'SEWT' computation with state s,
-- to another with state s. The current state and result of the given
-- computation is given to a mapping function that must produce the next
-- computation. The initial state must also be passed as the last parameter.
transitionState :: (Monad m, Monoid w) => ((a, s) -> SEWT t e w m a)
                -> m (SEWT s e w m a) -> s -> m (SEWT t e w m a)
transitionState _trans _comp _init = do
    (res, logs) <- _comp >>= flip runSEWT _init
    return $ do tell logs 
                case res of Left  fail -> throwError fail
                            Right succ -> _trans succ

-- 'withState': behaves like 'transitionState' but ignores the state of
-- the first computation.
withState :: (Monad m, Monoid w)
          => m (SEWT s e w m a) -> s -> m (SEWT t e w m a)
withState = transitionState $ return . fst

But is there perhaps a more elegant and general way to move from one state type to another?

I'm interested both in solutions where the second computation is not dependent on the final state (only the result) of the first computation, and one where it is.

Edit1: Improved transition functions:

transSEWT :: Functor m => (((a, y), x) -> (a, y)) -> SEWT x e w m a -> x -> SEWT y e w m a
transSEWT f x_c x_i = SEWT $ StateT $ \y_i -> ExceptT . WriterT $
    first ((\(a, x_f) -> f ((a, y_i), x_f)) <$>) <$> runSEWT x_c x_i

changeSEWT :: Functor m => SEWT x e w m a -> x -> SEWT y e w m a
changeSEWT = transSEWT fst

transS :: Monad m => (((a, y), x) -> (a, y)) -> StateT x m a -> x -> StateT y m a
transS f x_c x_i = StateT $ \y_i -> do (a, x_f) <- runStateT x_c x_i
                                       return $ f ((a, y_i), x_f)

changeS :: Monad m => StateT x m a -> x -> StateT y m a
changeS = transS fst

Answer 1

Your idea can be implemented with the indexed state monad.

newtype IState i o a = IState { runIState :: i -> (o, a) }

A value of type IState i o a is a stateful computation which returns a value of type a, transforming the type of the implicit state from i to o in the process. Contrast this with the regular State monad, which doesn't allow you to change the type of its state:

type State s = IState s s

Sequencing indexed state monads should ensure that the inputs and outputs line up. The output type of one computation is the input of the next. Enter Atkey's parameterised monad (now more commonly known as the indexed monad), the class of monad-like things describing a path through a directed graph.

class IMonad m where
    ireturn :: a -> m i i a
    (>>>=) :: m i j a -> (a -> m j k b) -> m i k b

(>>>) :: IMonad m => m i j a -> m j k b -> m i k b
mx >>> my = mx >>>= const my

Binding an indexed monad is like playing dominoes: if you have a way to get from i to j and a way to get from j to k, >>>= will glue your dominoes together into a bigger computation which goes from i to k. McBride describes a more powerful version of this indexed monad in Kleisli Arrows of Outrageous Fortune, but this one is quite enough for our purposes.

As I described above, domino-like sequencing is just what's needed for the indexed state monad, which requires alignment of inputs and outputs.

instance IMonad IState where
    ireturn x = IState $ \s -> (s, x)
    IState f >>>= g = IState $ \i -> let (o, x) = f i
                                     in runIState (g x) o

Retrieving a value from an indexed state monad doesn't change the type of the state.

get :: IState s s s
get = IState $ \s -> (s, s)

Putting a value into the indexed state monad discards the old state. This means the type of the input state can be anything you like.

put :: s -> IState i s ()
put x = IState $ \_ -> (x, ())

Note that all the code to work with IState is exactly the same as State! It's just the types which have got smarter.

Here's a simple IState computation which expects a state of type Int, changes the state to a String, and returns a Boolean answer. All of this is statically checked.

myStateComputation :: IState Int String Bool
myStateComputation =
    -- poor man's do notation. You could use RebindableSyntax
    get              >>>= \s ->
    put (show s)     >>>
    ireturn (s > 5)

main = print $ runIState myStateComputation 3
-- ("3", False)