{-# LANGUAGE Arrows #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE TupleSections #-}

{-
I implemented State and Reader arrows (actually as Arrow Transformers) for
educational purposes in order to understand arrows better.

Never use this code in a production environment. State and Reader arrows
have already been defined. I reimplemented them purely for educational
purposes.
-}

module Educational where

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

newtype Identity b c = Identity { runIdentity :: b -> c }

instance Category Identity where
    id = arr id
    (Identity c) . (Identity b) = Identity $ c . b

instance Arrow Identity where
    arr = Identity
    first (Identity b) = Identity $ Bi.first b
    second (Identity b) = Identity $ Bi.second b

class ArrowTrans t where
    lift :: (Arrow a) => a b c -> t a b c

newtype StateT s a b c = StateT { runStateT :: a (b, s) (c, s) }

instance Category a => Category (StateT s a) where
    id = StateT id
    (StateT c) . (StateT b) = StateT $ c . b
    
instance Arrow a => Arrow (StateT s a) where
   arr = StateT . first . arr 
   first (StateT f) = StateT $ arr (\((b, d), s) -> ((b, s), d)) >>> first f >>> arr (\((c, s), d) -> ((c, d), s))
   second (StateT f) = StateT $ arr (\((d, b), s) -> (d, (b, s))) >>> second f >>> arr (\(d, (c, s)) -> ((d, c), s))

class Arrow a => ArrowState s a | a -> s where
    gets :: a () s
    puts :: a s s 
    modify :: (s -> s) -> a () s 

instance Arrow a => ArrowState s (StateT s a) where
    gets = StateT $ arr $ \(_, s) -> (s, s)
    puts = StateT $ arr $ \(s, _) -> (s, s)
    modify t = StateT $ arr $ \(_, s) -> let s' = t s in (s', s')

instance ArrowTrans (StateT s) where
    lift = StateT . first

type State s b c = StateT s Identity b c

runState = runIdentity . runStateT

newtype ReaderT r a b c = ReaderT { runReaderT :: r -> a b c }

instance Category a => Category (ReaderT r a) where
    id = ReaderT $ const id 
    (ReaderT c) . (ReaderT b) = ReaderT $ \r -> c r . b r

instance Arrow a => Arrow (ReaderT r a) where
    arr = ReaderT . const . arr 
    first (ReaderT f) = ReaderT $ first . f 
    second (ReaderT f) = ReaderT $ second . f

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

instance Arrow a => ArrowReader r (ReaderT r a) where
    ask = ReaderT $ arr . const 

instance ArrowTrans (ReaderT r) where
    lift = ReaderT . const

type Reader r b c = ReaderT r Identity b c

runReader :: Reader r b c -> r -> b -> c 
runReader a r = runIdentity (runReaderT a r) 

zoo :: ArrowReader Int a => a Int Int
zoo = arr ((),) >>> first ask >>> arr (uncurry (*))

foo :: ArrowState Int a => a Int Int
foo = arr ((),) >>> first (modify (+3)) >>> second (arr (*2)) >>> arr (uncurry (+)) 

someFunc :: IO ()
someFunc = print $ runState foo (7, 9)