[Haskell] Simple variable-state `monad'

[oleg at pobox.com]

The familiar State monad lets us represent computation with a state
that can be queried and updated during the computation. The state must
have the same type during the entire computation however. One
sometimes wants to express a computation where not only the value but
the type of the state can be updated as well -- while maintaining
static typing.

David Roundy has recently suggested a similar `restricted monad'
> class WitnessMonad wm where
>    (>>=) :: wm w w' a -> (a -> wm w' w'' b) -> wm w w'' b
>    (>>)   :: wm w w' a -> wm w' w'' b -> wm w w'' b
>    return :: a -> wm w w' a
>    fail   :: String -> wm w w' a
>
> which obviously just adds a couple of phantom types indicating the monad
> state prior to and after each operation.  I actually rather like this
> redefined monad, as it seems like it'd be useful in lots of places (e.g. a
> restricted IO monad that disallows use of a file handle after it's
> closed).

To nitpick, such a `restricted monad' no longer can be called `monad'
because the type of the bind operation does not satisfy the signature
m a -> (a -> m b) -> m b. Proper naming is important so we can
understand what kind of animal we're dealing with and so can chose the
right way to domesticate it.

The problem indeed shows up in many contexts, like ensuring that a
file cannot be used after it is closed, or that the lock is properly
released. The most common solution to these problems (expressing
control-flow proiperties) involves using monad in a contra-variant
position. One recent example

	http://pobox.com/~oleg/ftp/Haskell/yield.hs 

is ensuring a locking property (a lock can be unlocked only when it is
held; one can acquire the same lock several times but must release
that many times; no `yield' may be invoked while holding a
lock). However, this style is a bit ugly. One wonders if a more direct
expression of this `variable style monad' is possible. It seems it is,
as the following code demonstrates. It runs with GHC 6.4 and Hugs.

{-# OPTIONS -fglasgow-exts #-}

-- Variable state monad

module EM1 where

import Control.Monad.State

class Monadish m where
    gret :: a -> m p p a
    gbind :: m p q a -> (a -> m q r b) -> m p r b

-- Inject regular monads to be monadish things too
newtype MW m p q a = MW{ unMW:: m a }

instance Monad m => Monadish (MW m) where
    gret = MW . return
    gbind (MW m) f = MW (m >>= unMW . f)

-- As a warm-up, we write the regular State computation, with the same 
-- type of state throughout

-- First, use the regular Monad.State
test1 = runState c (0::Int) where
         c = do
             v <- get
             put (succ v)
             return v
-- (0,1)

-- Now, wrap in MW
test2 = runState (unMW c) (0::Int) where
         c = gget `gbind` (
               \v -> gput (succ v) `gbind` (
                 \_ -> gret v))
         gget :: (MonadState s m) => MW m s s s
         gget = MW get
         gput :: (MonadState s m) => s -> MW m s s ()
         gput = MW . put
-- (0,1)


-- Introduce the variable-type state

newtype VST m si so v = VST{runVST:: si -> m (so,v)}


instance Monad m => Monadish (VST m) where
  gret x = VST (\si -> return (si,x))
  gbind (VST m) f = VST (\si -> m si >>= (\ (sm,x) -> runVST (f x) sm))

vsget :: Monad m => VST m si si si
vsget = VST (\si -> return (si,si))
vsput :: Monad m => so -> VST m si so ()
vsput x = VST (\si -> return (x,()))


-- Repeat test2 via VST: the type of the state is the same
test3 = runVST c (0::Int) >>= print where
         c = vsget `gbind` (
               \v -> vsput (succ v) `gbind` (
                 \_ -> gret v))
-- (1,0)

-- Now, we vary the type of the state, from Int to a Char
test4 = runVST c (0::Int) >>= print where
         c = vsget `gbind` (
               \v -> vsput ((toEnum (65+v))::Char) `gbind` (
                 \_ -> vsget `gbind` (
                   \v' -> gret (v,v'))))
-- ('A',(0,'A'))

{-
-- The following is an error: 
    Couldn't match `Char' against `Int'
      Expected type: VST m Char r b
      Inferred type: VST m Int Int Bool

test4' = runVST c (0::Int) >>= print where
         c = vsget `gbind` (
               \v -> vsput ((toEnum (65+v))::Char) `gbind` (
                 \_ -> vsget `gbind` (
                   \v' -> gret (v==v'))))
-}


-- Try polymorphic recursion, over the state.
-- crec1 invokes itself, and changes the type of the state from
-- some si to Bool.
crec1 :: (Enum si, Monad m) => VST m si si Int
crec1 = vsget `gbind` (\s1 -> case fromEnum s1 of
                      0 -> gret 0
                      1 -> vsput (pred s1) `gbind` (\_ -> gret 1)
                      _ -> vsput True `gbind` (\_ -> 
                             crec1 `gbind` (\v ->
                             (vsput s1 `gbind` -- restore state type to si
                             (\_ -> gret (v + 10))))))

test5 = runVST crec1 'a' >>= print
-- ('a',11)