State Monad

module Main where

--             = Reader s (Writer s a)
data State s a = State (s -> (a, s))

runState :: State s a -> s -> (a, s)
runState (State f) s = f s

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

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

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

instance Applicative (State s) where
  -- pure :: a -> State s a
  pure x = State (\s -> (x, s))

  -- (<*>) :: State s (a -> b) -> State s a -> State s b
  sab <*> sa = State (\s -> let (f, s1) = runState sab s
                                (a, s2) = runState sa  s1
                            in  (f a, s2))

instance Monad (State s) where
  -- (>>=) :: State s a -> (a -> State s b) -> State s b
  sa >>= f = State (\s -> let (a, s1) = runState sa s
                              sb      = f a
                          in  runState sb s1)

-- Exemplo
data Tree a = Leaf a | Node (Tree a) (Tree a) deriving Show

inc :: State Int Int
inc = State (\n -> (n, n+1))

rlabel :: Tree a -> Int -> (Tree Int, Int)
rlabel (Leaf _) n   = (Leaf n, n+1)
rlabel (Node l r) n = let (l', n1) = rlabel l n
                          (r', n2) = rlabel r n1
                      in  (Node l' r', n2)

alabel :: Tree a -> State Int (Tree Int)
alabel (Leaf _) = pure Leaf <*> inc
alabel (Node l r) = pure Node <*> alabel l <*> alabel r

mlabel :: Tree a -> State Int (Tree Int)
mlabel (Leaf _) = do n <- inc
                     return (Leaf n)
mlabel (Node l r) = do l' <- mlabel l
                       r' <- mlabel r -- e se eu trocar a ordem?
                       return $ Node l' r'

t = Node (Node (Node (Leaf 'a') (Leaf 'b')) (Leaf 'c')) (Node (Leaf 'd') (Leaf 'e'))

main = do
  print $ rlabel t 1
  print $ runState (alabel t) 1
  print $ runState (mlabel t) 1