Free State

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
{-# LANGUAGE DeriveFunctor #-}

import Control.Monad.Trans.Free
import Control.Monad.Trans.Class (lift)
import Data.Composition ((.:))
import Control.Monad.Identity (Identity, runIdentity)


data StateF s x
  = Get (s -> x)
  | Put s x
  deriving Functor

type StateT s = FreeT (StateF s)
type State s = StateT s Identity

get :: Monad m => StateT s m s
get = liftF $ Get id

put :: Monad m => s -> StateT s m ()
put s = liftF $ Put s ()

modify :: Monad m => (s -> s) -> StateT s m ()
modify f = get >>= put . f

stateT :: Monad m => (s -> m (a, s)) -> StateT s m a
stateT f = do
  s <- get
  ~(a, s') <- lift (f s)
  put s'
  return a

state :: Monad m => (s -> (a, s)) -> StateT s m a
state f = stateT (return . f)

runStateT :: Monad m => StateT s m a -> s -> m (a, s)
runStateT m s = do
  x <- runFreeT m
  case x of
    Pure a -> return (a, s)
    Free (Get f) -> runStateT (f s) s
    Free (Put s' next) -> runStateT next s'

runState :: State s a -> s -> (a, s)
runState = runIdentity .: runStateT