{-# LANGUAGE FlexibleInstances #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Control.Arrow.Transformer.State
-- Copyright   :  (c) Ross Paterson 2003
-- License     :  BSD-style (see the LICENSE file in the distribution)
--
-- Maintainer  :  ross@soi.city.ac.uk
-- Stability   :  experimental
-- Portability :  non-portable (multi-parameter type classes)
--
-- An arrow transformer that adds a modifiable state,
-- based of section 9 of /Generalising Monads to Arrows/, by John Hughes,
-- /Science of Computer Programming/ 37:67-111, May 2000.

module Control.Arrow.Transformer.State(
		StateArrow(StateArrow),
		runState,
		ArrowAddState(..),
	) where

import Control.Arrow.Internals
import Control.Arrow.Operations
import Control.Arrow.Transformer

import Control.Applicative
import Control.Arrow
import Control.Category
import Data.Monoid

import Prelude hiding (id,(.))

-- | An arrow type that augments an existing arrow with a modifiable
-- state.  The 'ArrowState' class contains the operations on this state.

newtype StateArrow s a b c = StateArrow (a (b, s) (c, s))

swapsnd :: ((a, b), c) -> ((a, c), b)
swapsnd ~(~(x, y), z) = ((x, z), y)

instance Category a => Category (StateArrow s a) where
	id = StateArrow id
	StateArrow f . StateArrow g = StateArrow (f . g)

instance Arrow a => Arrow (StateArrow s a) where
	arr f = StateArrow (arr (\(x, s) -> (f x, s)))
	first (StateArrow f) =
		StateArrow (arr swapsnd >>> first f >>> arr swapsnd)

instance Arrow a => ArrowTransformer (StateArrow s) a where
	lift f = StateArrow (first f)

-- | Encapsulation of a state-using computation, exposing the initial
-- and final states.
--
-- Typical usage in arrow notation:
--
-- >	proc p -> do
-- >		...
-- >		(result, final_state) <- (|runState cmd|) init_state

runState :: Arrow a => StateArrow s a e b -> a (e,s) (b,s)
runState (StateArrow f) = f

-- operations

instance Arrow a => ArrowState s (StateArrow s a) where
	fetch = StateArrow (arr (\(_, s) -> (s, s)))
	store = StateArrow (arr (\(s, _) -> ((), s)))

instance Arrow a => ArrowAddState s (StateArrow s a) a where
	liftState = lift
	elimState = runState

-- The following promotions follow directly from the arrow transformer.

instance ArrowZero a => ArrowZero (StateArrow s a) where
	zeroArrow = StateArrow zeroArrow

instance ArrowCircuit a => ArrowCircuit (StateArrow s a) where
	delay x = lift (delay x)

instance ArrowError ex a => ArrowError ex (StateArrow s a) where
	raise = lift raise
	handle (StateArrow f) (StateArrow h) =
		StateArrow (handle f (arr swapsnd >>> h))
	tryInUnless (StateArrow f) (StateArrow s) (StateArrow h) =
		StateArrow (tryInUnless f (arr new_state >>> s) (arr swapsnd >>> h))
		where	new_state ((b,_),(c,s')) = ((b,c),s')
	newError (StateArrow f) = StateArrow (newError f &&& arr snd >>> arr h)
		where	h (Left ex, s) = (Left ex, s)
			h (Right (c, s'), _) = (Right c, s')

-- Note that in each case the error handler gets the original state.

instance ArrowReader r a => ArrowReader r (StateArrow s a) where
	readState = lift readState
	newReader (StateArrow f) = StateArrow (arr swapsnd >>> newReader f)

instance ArrowWriter w a => ArrowWriter w (StateArrow s a) where
	write = lift write
	newWriter (StateArrow f) = StateArrow (newWriter f >>> arr swapsnd)

-- liftings of standard classes

instance ArrowChoice a => ArrowChoice (StateArrow s a) where
	left (StateArrow f) = StateArrow (arr distr >>> left f >>> arr undistr)
		where	distr (Left y, s) = Left (y, s)
			distr (Right z, s) = Right (z, s)
			undistr (Left (y, s)) = (Left y, s)
			undistr (Right (z, s)) = (Right z, s)

instance ArrowApply a => ArrowApply (StateArrow s a) where
	app = StateArrow (arr (\((StateArrow f, x), s) -> (f, (x, s))) >>> app)

instance ArrowLoop a => ArrowLoop (StateArrow s a) where
	loop (StateArrow f) =
		StateArrow (loop (arr swapsnd >>> f >>> arr swapsnd))

instance ArrowPlus a => ArrowPlus (StateArrow s a) where
	StateArrow f <+> StateArrow g = StateArrow (f <+> g)

-- Other instances

instance Arrow a => Functor (StateArrow s a b) where
	fmap f g = g >>> arr f

instance Arrow a => Applicative (StateArrow s a b) where
	pure x = arr (const x)
	f <*> g = f &&& g >>> arr (uncurry id)

instance ArrowPlus a => Alternative (StateArrow s a b) where
	empty = zeroArrow
	f <|> g = f <+> g

instance ArrowPlus a => Monoid (StateArrow s a b c) where
	mempty = zeroArrow
	mappend f g = f <+> g

-- promotions

instance ArrowAddReader r a a' =>
		ArrowAddReader r (StateArrow s a) (StateArrow s a') where
	liftReader (StateArrow f) = StateArrow (liftReader f)
	elimReader (StateArrow f) = StateArrow (arr swapsnd >>> elimReader f)

instance ArrowAddWriter w a a' =>
		ArrowAddWriter w (StateArrow s a) (StateArrow s a') where
	liftWriter (StateArrow f) = StateArrow (liftWriter f)
	elimWriter (StateArrow f) = StateArrow (elimWriter f >>> arr swapsnd)

instance ArrowAddError ex a a' =>
		ArrowAddError ex (StateArrow s a) (StateArrow s a') where
	liftError (StateArrow f) = StateArrow (liftError f)
	elimError (StateArrow f) (StateArrow h) =
		StateArrow (elimError f (arr swapsnd >>> h))