-- | Motivation and illustration of the writer monad.
module Logging where

import Control.Monad (ap,liftM,liftM2)
import Data.Sequence (Seq,(|>),singleton)


-- * Syntax

-- | An arithmetic expression language that can log intermediate results.
data Expr
   = Lit Int           -- ^ literal integer
   | Neg Expr          -- ^ integer negation
   | Bin Op Expr Expr  -- ^ binary operator
   | Log Expr          -- ^ log the result and return it
  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 (Log (Lit 4)) (Log (Lit 5))        -- ~> ([4,5],9)
e2 = Log e1                                 -- ~> ([4,5,9],9)
e3 = mul (Log (Neg (Log (Lit 2)))) (Lit 3)  -- ~> ([2,-2],-6) 
e4 = add e1 e3                              -- ~> ([4,5,2,-2],3)


-- * Semantics

-- ** Recursive interpreter

-- | The log of intermediate results.
type Log = Seq Int

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

-- | Evaluate an expression.
eval :: Expr -> (Log, Int)
eval (Lit i)     = (mempty, i)
eval (Neg e)     = let (w, i) = eval e
                   in (w, negate i)
eval (Bin o l r) = let (wl, i) = eval l
                       (wr, j) = eval r
                   in (wl <> wr, op o i j)
eval (Log e)     = let (w, i) = eval e
                   in (w |> i, i)



-- ** Monadic interpreter

-- *** Define the monad

-- | A computation that produces output as a side effect.
newtype Writer w a = W (w, a)
  deriving (Eq,Show)

instance Functor (Writer w) where
  -- fmap :: (a -> b) -> Writer w a -> Writer w b
  fmap f (W (s, x)) = W (s, f x)

instance Monoid w => Applicative (Writer w) where
  
  -- pure :: a -> Writer w a
  pure x = W (mempty, x)

  -- (<*>) :: Writer w (a -> b) -> Writer w a -> Writer w b
  W (s1, f) <*> W (s2, x) = W (s1 <> s2, f x)

instance Monoid w => Monad (Writer w) where

  -- return :: a -> Writer w a
  return x = W (mempty, x)

  -- (>>=) :: Writer w a -> (a -> Writer w b) -> Writer w b
  W (s, x) >>= f = let W (s', y) = f x
                   in W (s <> s', y)


-- *** Define the "primitive" operations

-- | Run a writer action.
runWriter :: Writer w a -> (w, a)
runWriter (W p) = p

-- | Produce some output.
tell :: w -> Writer w ()
tell s = W (s, ())


-- *** Use the monad

-- | Monadic evaluation function.
evalM :: Expr -> Writer Log Int
evalM (Lit i)     = return i
evalM (Neg e)     = do i <- evalM e
                       return (negate i)   -- evalM e >>= \i -> return (negate i)
evalM (Bin o l r) = do i <- evalM l
                       j <- evalM r
                       return (op o i j)   -- evalM l >>= \i -> (evalM r >>= \j -> return (op o i j))
evalM (Log e)     = do i <- evalM e
                       tell (singleton i)
                       return i            -- evalM e >>= \i -> (tell (singleton i) >> return i)

-- evalM (Lit i)     = pure i
-- evalM (Neg e)     = negate <$> evalM e
-- evalM (Bin o l r) = op o <$> evalM l <*> evalM r
-- evalM (Log e)     = do i <- evalM e
--                        tell (singleton i)
--                        pure i

-- evalM (Neg e)     = fmap negate (evalM e)
-- evalM (Bin o l r) = liftA2 (op o) (evalM l) (evalM r)