state stuff

auto_state.hs

{-# LANGUAGE DeriveTraversable #-}

module AS where

import Control.Monad.Trans.State.Strict

data T a = L a | N (T a) (T a)
  deriving (Eq, Show, Functor, Foldable, Traversable)

labelTree :: T a -> (T (a, Int), Int)
labelTree t = runState (tl t) 0 where
  tl = traverse (\c -> state (\n -> ((c, n), n + 1)))

manual_state.hs

{-# LANGUAGE InstanceSigs #-}

module MS where

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

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

instance Applicative (State s) where
  pure :: a -> State s a
  pure a = State $ \s -> (a, s)
  (<*>) :: State s (a -> b) -> State s a -> State s b
  (State stf) <*> (State sta) = State $ \s ->
   let
    (f, s')  = stf s
    (a, s'') = sta s'
   in
    (f a, s'')

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

state :: (s -> (a, s)) -> State s a
state = State

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

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

evalState :: State s a -> s -> a
evalState st = fst . runState st

execState :: State s a -> s -> s
execState st = snd . runState st

--

data Tree a = Leaf a | Node (Tree a) (Tree a) deriving (Eq, Show)

instance Functor Tree where
  fmap :: (a -> b) -> Tree a -> Tree b
  fmap f (Leaf x)   = Leaf (f x)
  fmap f (Node l r) = Node (fmap f l) (fmap f r)

instance Foldable Tree where
  foldMap :: Monoid b => (a -> b) -> Tree a -> b
  foldMap f (Leaf x)   = f x
  foldMap f (Node l r) = foldMap f l <> foldMap f r

instance Traversable Tree where
  traverse :: Applicative f => (a -> f b) -> Tree a -> f (Tree b)
  traverse f (Leaf x)   = Leaf <$> f x
  traverse f (Node l r) = Node <$> traverse f l <*> traverse f r

--

manualLabel :: Eq a => Tree a -> Int -> (Tree (a, Int), Int)
manualLabel (Leaf a  ) n = (Leaf (a, n), n + 1)
manualLabel (Node l r) n = (Node l' r', n'')
 where
  (l', n' ) = manualLabel l n
  (r', n'') = manualLabel r n'

appLabel :: Eq a => Tree a -> State Int (Tree (a, Int))
appLabel (Leaf a  ) = Leaf <$> state (\n -> ((a, n), n + 1))
appLabel (Node l r) = Node <$> appLabel l <*> appLabel r

monadLabel :: Eq a => Tree a -> State Int (Tree (a, Int))
monadLabel (Leaf a) = do
  n <- get
  put (n + 1)
  return (Leaf (a, n))
monadLabel (Node l r) = do
  l' <- monadLabel l
  r' <- monadLabel r
  return (Node l' r')

traverseLabel :: Tree a -> State Int (Tree (a, Int))
traverseLabel = traverse label
 where
  label :: a -> State Int (a, Int)
  label c = do
    n <- get
    put (n + 1)
    return (c, n)

--

tree :: Tree Char
tree = Node (Node (Leaf 'a') (Leaf 'b')) (Node (Leaf 'c') (Leaf 'd'))

runManualLabel :: (Tree (Char, Int), Int)
runManualLabel = runState (state (manualLabel tree)) 0

runAppLabel :: (Tree (Char, Int), Int)
runAppLabel = runState (appLabel tree) 0

runMonadLabel :: (Tree (Char, Int), Int)
runMonadLabel = runState (monadLabel tree) 0

runTraverseLabel :: (Tree (Char, Int), Int)
runTraverseLabel = runState (traverseLabel tree) 0