RWS.hs

{-# LANGUAGE FlexibleContexts, FlexibleInstances, MultiParamTypeClasses, TypeOperators, UndecidableInstances, DeriveFunctor, StandaloneDeriving, RankNTypes #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module RWS where

import Control.Applicative
import Control.Effect
import Control.Effect.Carrier
import Control.Effect.Reader
import Control.Effect.State.Internal
import Control.Effect.Sum
import Control.Effect.Writer

newtype RWSC r w s m a = RWSC { unRWSC :: r -> s -> w -> m (s, w, a) }
  deriving Functor

instance Monad m => Applicative (RWSC r w s m) where
  pure a = RWSC $ \ _ s w -> pure (s, w, a)
  {-# INLINE pure #-}

  RWSC mf <*> RWSC mx = RWSC $ \ r s w -> do
    (s', w', f)   <- mf r s w
    (s'', w'', x) <- mx r s' w'
    pure (s'', w'', f x)
  {-# INLINE (<*>) #-}

instance Monad m => Monad (RWSC r w s m) where
  m >>= k = RWSC $ \ r s w -> do
    (s', w', a) <- unRWSC m r s w
    unRWSC (k a) r s' w'
  {-# INLINE (>>=) #-}

instance (Monad m, Alternative m) => Alternative (RWSC r w s m) where
  empty = RWSC $ \ _ _ _ -> empty
  {-# INLINE empty #-}

  RWSC m <|> RWSC n = RWSC $ \ r s w -> m r s w <|> n r s w
  {-# INLINE (<|>) #-}

deriving instance Functor ((,,) a b)

runRWS :: r -> w -> s -> RWSC r w s m a -> m (s, w, a)
runRWS r w s x = unRWSC x r s w

instance (Monoid w, Carrier sig m, Effect sig) => Carrier (Reader r :+: Writer w :+: State s :+: sig) (RWSC r w s m) where
  -- eff (L (Cull m k))     = CullC (local (const True) (runCullC m)) >>= k
  -- eff (R (L Empty))      = empty
  eff (L (Ask k)) = RWSC $ \r w s -> unRWSC (k r) r w s
  eff (L (Local f m k)) = RWSC (\r w s -> unRWSC m (f r) w s) >>= k
  eff (R (L (Tell w' k))) = RWSC (\r s w -> let wt = w `mappend` w' in wt `seq` unRWSC k r s wt )
  eff (R (L (Listen m k))) = RWSC $ \r s w -> do
    (s', w', a) <- unRWSC m r s w
    let wt = wt `mappend` w'
    wt `seq` unRWSC (k w' a) r s' wt
  eff (R (L (Censor f m k))) = RWSC $ \r s w -> do
    (s', w', a) <- unRWSC m r s w
    let wt = wt `mappend` f w'
    wt `seq` unRWSC (k a) r s' wt
  eff (R (R (L (Get k)))) = RWSC $ \r s w -> unRWSC (k s) r s w
  eff (R (R (L (Put s' k)))) = RWSC $ \r _ w -> unRWSC k r s' w
  eff (R (R (R other))) = RWSC $ \r w s -> eff (handle (w, s, ()) (\(s', w', go) -> unRWSC go r s' w') other)
  {-# INLINE eff #-}