-- | Motivation and illustration of the state monad.
module State where

import Control.Monad (ap,liftM)


-- * Syntax

-- | A language for a calculator with a single register of memory.
data Expr
   = Lit Int           -- ^ literal integer
   | Neg Expr          -- ^ integer negation
   | Bin Op Expr Expr  -- ^ binary operator
   | Set Expr Expr     -- ^ set register to result of LHS, evaluate RHS
   | Save Expr         -- ^ save result to register (and return it)
   | Load              -- ^ load from register
  deriving (Eq,Show)

-- | Binary operators.
data Op = Add | Mul
  deriving (Eq,Show)

-- | Addition smart constructor.
add :: Expr -> Expr -> Expr
add = Bin Add

-- | Addition smart constructor.
mul :: Expr -> Expr -> Expr
mul = Bin Mul

-- | Example expressions.
e1 = add (Save (Lit 2)) Load                   -- ~>  4   2
e2 = Set (Lit 3) (mul Load (Neg Load))         -- ~> -9   3
e3 = Set (Neg (add e1 e2))                     -- ~>  5   5
         (add (Save (add (Lit 1) Load)) Load)  -- ~> 12   6


-- * Semantics

-- ** Recursive interpreter

-- | The register value.
type Reg = Int

-- | Function corresponding to each binary operator.
op :: Op -> Int -> Int -> Int
op Add = (+)
op Mul = (*)

-- | Evaluate an expression.
eval :: Expr -> Reg -> (Int, Reg)
eval (Lit n)     s = (n, s)
eval (Neg e)     s = let (i, s') = eval e s
                     in (negate i, s')
eval (Bin o l r) s = let (i, s')  = eval l s
                         (j, s'') = eval r s'
                     in (op o i j, s'')
eval (Set l r)   s = let (i, _)  = eval l s
                     in eval r i
eval (Save e)    s = let (i, _) = eval e s
                     in (i, i)
eval Load        s = (s, s)


-- ** Monadic interpreter

-- | A stateful computation.
newtype State s a = S (s -> (a, s))

instance Functor (State s) where
  fmap = liftM

instance Applicative (State s) where
  pure  = return
  (<*>) = ap

instance Monad (State s) where

  -- return :: a -> State s a
  return x = S (\s -> (x, s))

  -- (>>=) :: State s a -> (a -> State s b) -> State s b
  S f >>= g = S (\s -> let (x, s') = f s
                           S h = g x
                       in h s')

  -- f :: s -> (a, s)
  -- g :: a -> State s b
  -- h :: s -> (b, s)
  -- x :: a
  -- s, s' :: s

-- | Run the computation with a given initial state.
runState :: State s a -> s -> (a, s)
runState (S f) s = f s

-- | Set the curent state.
put :: s -> State s ()
put s = S (\_ -> ((), s))

-- | Get the current state.
get :: State s s
get = S (\s -> (s, s))


-- Border between library and client --

-- | Monadic evaluation function.
evalM :: Expr -> State Reg Int
evalM (Lit n)     = return n
evalM (Neg e)     = do i <- evalM e
                       return (negate i)
evalM (Bin o l r) = do i <- evalM l
                       j <- evalM r
                       return (op o i j)
evalM (Set l r)   = do i <- evalM l
                       put i
                       evalM r
evalM (Save e)    = do i <- evalM e
                       put i
                       return i
evalM Load        = get

-- evalM (Lit n)     = return n
-- evalM (Neg e)     = evalM e >>= \i -> return (negate i)
-- evalM (Bin o l r) = evalM l >>= \i -> evalM r >>= \j -> return (op o i j)
-- evalM (Set l r)   = evalM l >>= put >> evalM r
-- evalM (Save e)    = evalM e >>= \i -> put i >> return i
-- evalM Load        = get

-- evalM (Lit n)     = return n
-- evalM (Neg e)     = negate <$> evalM e
-- evalM (Bin o l r) = op o <$> evalM l <*> evalM r
-- evalM (Set l r)   = evalM l >>= put >> evalM r
-- evalM (Save e)    = evalM e >>= put >> get
-- evalM Load        = get

-- evalM (Neg e)     = fmap negate (evalM e)
-- evalM (Bin o l r) = liftA2 (op o) (evalM l) (evalM r)
-- evalM (Save e)    = do { i <- evalM e; put i; return i }

-- | Evaluate an expression with an initial state of 0.
eval' :: Expr -> (Int, Reg)
eval' e = runState (evalM e) 0