-- Imports {{{
-- vim: fdm=marker
{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses,
             FunctionalDependencies, KindSignatures #-}
module Lecture11 where

import Control.Monad.Reader.Class
import Control.Monad.Writer.Class
import Control.Monad ( liftM2 )
import Data.Monoid ( Sum ( getSum ) )
import Control.Applicative ( liftA2 )
import Data.Map ( Map )
import qualified Data.Map as Map
-- }}}

-- Type for formulae {{{
data Formula = Var String | And Formula Formula
             | Not Formula | Or Formula Formula
             deriving (Eq, Ord, Show)
type Valuation = Map String Bool
-- }}}

-- Our combined Reader + Writer monad {{{
newtype RW r w a = RW { runRW :: r -> (w, a) }
instance Functor (RW r w) where
    fmap f rwa = RW $ \r -> f <$> runRW rwa r
instance Monoid w => Applicative (RW r w) where
    liftA2 = liftM2
    pure x = RW $ const (mempty, x)
instance Monoid w => Monad (RW r w) where
    mx >>= f = RW $ \r -> let (w1, x) = runRW mx r
                              (w2, y) = runRW (f x) r
                          in (w1 <> w2, y)
instance Monoid w => MonadReader r (RW r w) where
    ask = RW $ \r -> (mempty, r)
    local f rwa = RW $ \r -> runRW rwa (f r)

instance Monoid w => MonadWriter w (RW r w) where
    tell x = RW $ const (x, ())
    listen rwa = RW $ \r -> let (w, a) = runRW rwa r in (w, (a, w))
    pass rwaf = RW $ \r -> let (w, (a, f)) = runRW rwaf r in (f w, a)
-- }}}

-- Evaluator using our RW monad {{{
eval :: Formula -> RW Valuation (Sum Integer) Bool
eval (Var v)   = do tell 1
                    asks $ Map.findWithDefault False v
eval (And x y) = do lhs <- eval x
                    tell 1
                    if not lhs then pure False
                               else eval y
eval (Or  x y) = tell 1 >> liftA2 (||) (eval x) (eval y)
eval (Not x)   = tell 1 >> not <$> eval x

countSteps :: Formula -> Valuation -> Integer
countSteps f vs = getSum . fst $ runRW (eval f) vs
-- }}}