Haskell ArrowReader & ArrowState

{-# LANGUAGE Arrows #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}

module Arr where

{-

As an exercise in learning arrows, I implemented a reader arrow called EA
and a state arrow called S.

-}

import Prelude hiding ((.), id)
import Control.Category
import Control.Arrow
import qualified Data.Bifunctor as BiF

newtype EA r a b = EA { runEA :: r -> a -> b }

instance Arrow (EA r) where
  arr f = EA $ const f
  first (EA f) = EA $ BiF.first . f
  second (EA f) = EA $ BiF.second . f

instance Category (EA r) where
  id = arr id
  (EA b) . (EA a) = EA $ \r -> b r . a r

class Arrow a => ArrowReader r a | a -> r where
  ask :: a () r

instance ArrowReader r (EA r) where
  ask = EA const

newtype S s a b = S { runS :: s -> a -> (b, s) }

evalS :: S s a b -> s -> a -> b
evalS z s a = fst $ runS z s a

execS :: S s a b -> s -> a -> s
execS z s a = snd $ runS z s a

instance Arrow (S s) where
  arr f = S $ \s a -> (f a, s)
  first (S f) = S $ \s (a, b) -> let (a', s') = f s a in ((a', b), s')
  second (S f) = S $ \s (a, b) -> let (b', s') = f s b in ((a, b'), s')

instance Category (S s) where
  id = arr id
  (S b) . (S a) = S $ \s x -> let (a', s') = a s x in b s' a'

class Arrow a => ArrowState s a | a -> s where
  get :: a b s
  put :: a s ()

modify :: ArrowState s a => a (s -> s) ()
modify = proc f -> do
  s <- get -< ()
  put -< f s

instance ArrowState s (S s) where
  get = S $ \s _ -> (s, s)
  put = S $ \_ s -> ((), s)