State Monad

StateMonad.hs

newtype State s a = State { runState :: s -> (s, a) }

instance Functor (State s) where
  -- fmap :: (a -> b) -> State s a -> State s b
  fmap g (State f) = State (\s ->
    let (s', a) = f s
    in (s', g a))

instance Applicative (State s) where
  -- pure :: a -> (State s a)
  pure v = State (\s -> (s, v))
  -- State s (a -> b) -> State s a -> State s b
  (State sF) <*> (State sV) = State (\s ->
    let (s', f)  = sF s
        (s'', v) = sV s'
    in (s'', f v))

instance Monad (State s) where
  return = pure
  -- (>>=) :: State s a -> (a -> State s b) -> State s b
  (>>=) (State sA) f = State (\s ->
    let (s', a) = sA s
    in runState (f a) s')

get :: State s s
get = State (\s -> (s, s))

put :: s -> State s ()
put s = State (\_ -> (s, ()))

modify :: (s -> s) -> State s ()
modify f = get >>= put . f