[Haskell-cafe] AT solution: rebinding >>= for restricted monads

[Malcolm Wallace]

"Iavor Diatchki" <iavor.diatchki at gmail.com> wrote:

> > class IxMonad m where
> >    (>>>=) :: m i j a -> (a -> m j k b) -> m i k b
> >    ret    :: a -> m i i a
> 
> And just for fun we can define another indexed monad:
> state that supports "strong updates" (i.e., the type
> of the state can change as you compute):
> 
> > newtype State i j a = S { unS :: i -> (a,j) }

This reminded me very much of the state monad (actually more general
than a monad) used internally in the nhc98/yhc compiler.  Attached
below.  Niklas Rojemo came up with this formulation in 1993.

Regards,
    Malcolm

----
module State where

 -- Use >>> and >>>= instead of >> and >>= to avoid problems with 1.3 prelude
infixl 5 >>>,>>>=,=>>>
infixr 4 >=>

type State0  d s   s' = d -> s -> s'
type State   d s u s' = d -> s -> (u,s')

-- 'unit' is what is now known as 'return' in monad-speak
unitS :: u -> State d s u s
unitS u = (\d s -> (u,s))

unitS0 :: State0 d s   s
unitS0 = (\d s -> s)

-- There are four 'bind'-like operators, >>>, >>>=, =>>>, and >=>

(>>>) :: State0 d a b -> State0 d b c -> State0 d a c
f >>> g = \d s -> g d (f d s)

(=>>>) :: State d s (a->b) s' -> State d s' a s'' -> State d s b s''
f =>>> g = \d s -> case f d s of (h,s) ->
                       case g d s of (x,s) ->
                           let hx = h x in seq hx (hx,s)
-- f =>>> g = \d s -> case f d s of (h,s) -> case g d s of (x,s) -> (h x,s)

(>>>=) :: State d s a s' -> (a -> State0 d s' s'') -> State0 d s s''
f >>>= g = \d s -> case f d s of (x,s) -> g x d s

(>=>) :: State d s d' s' -> State0 d' s' s'' -> State0 d s s''
f >=> g = \d s -> case f d s of (d,s) -> (g d s)

mapS :: (a->State d s b s) -> [a] -> State d s [b] s
mapS f []     = unitS []
mapS f (x:xs) = unitS (:) =>>> f x =>>> mapS f xs

mapS0 :: (a->State0 d s s) -> [a] -> State0 d s s
mapS0 f []     = unitS0
mapS0 f (x:xs) = f x >>> mapS0 f xs

zipWithS :: (a -> b -> State d s c s) -> [a] -> [b] -> State d s [c] s
zipWithS f [] [] = unitS []
zipWithS f (x:xs) (y:ys) = unitS (:) =>>> f x y =>>> zipWithS f xs ys
zipWithS _ _ _ = error "zipWithS: lists of different lengths"

----