content: {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE UndecidableInstances #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE RoleAnnotations #-}
  {-# LANGUAGE DataKinds #-}
  module Control.Arrow.Transformer.State(StateT(..),evalStateT,execStateT,withStateT) where
  
  import           Prelude hiding (id,(.),lookup,read,fail)
  
  import           Control.Category
  import           Control.Arrow
  import           Control.Arrow.Cont
  import           Control.Arrow.Const
  import           Control.Arrow.Primitive
  import           Control.Arrow.Reader as Reader
  import           Control.Arrow.State as State
  import           Control.Arrow.Trans
  import           Control.Arrow.Writer
  
  import           Unsafe.Coerce
  import           Data.Monoidal
  import           Data.Profunctor hiding (Strong(..))
  import           Data.Profunctor.Unsafe
  
  -- Due to "Generalising Monads to Arrows", by John Hughes, in Science of Computer Programming 37.
  -- | Arrow transformer that adds mutable state to an arrow computation.
  newtype StateT s c x y = StateT { runStateT :: c (s,x) (s,y) }
  
  evalStateT :: Profunctor c => StateT s c x y -> c (s,x) y
  evalStateT f = rmap snd $ runStateT f
  {-# INLINE evalStateT #-}
  
  execStateT :: Profunctor c => StateT s c x y -> c (s,x) s
  execStateT f = rmap fst $ runStateT f
  {-# INLINE execStateT #-}
  
  withStateT :: Arrow c => StateT s' c x y -> StateT s c (s',x) (s',y)
  withStateT f = lift (second (unlift f))
  {-# INLINE withStateT #-}
  
  instance ArrowRun c => ArrowRun (StateT s c) where
    type Run (StateT s c) x y = Run c (s,x) (s,y)
  
  instance ArrowLift (StateT s c) where
    type Underlying (StateT s c) x y = c (s,x) (s,y)
  
  instance (ArrowPrimitive c) => ArrowPrimitive (StateT s c) where
    type PrimState (StateT s c) = PrimState c
  
  instance (Profunctor c) => Profunctor (StateT s c) where
    dimap f g h = lift $ dimap (second' f) (second' g) (unlift h)
    lmap f h = lift $ lmap (second' f) (unlift h)
    rmap g h = lift $ rmap (second' g) (unlift h)
    f .# _ = f `seq` unsafeCoerce f
    _ #. g = g `seq` unsafeCoerce g
    {-# INLINE dimap #-}
    {-# INLINE lmap #-}
    {-# INLINE rmap #-}
    {-# INLINE (.#) #-}
    {-# INLINE (#.) #-}
  
  instance ArrowTrans (StateT s) where
    lift' f = lift (second f)
    {-# INLINE lift' #-}
  
  instance (Arrow c) => Category (StateT s c) where
    id = lift id
    f . g = lift (unlift f . unlift g)
    {-# INLINE id #-}
    {-# INLINE (.) #-}
  
  instance (Arrow c, Profunctor c) => Arrow (StateT s c) where
    arr f = lift (arr (second' f))
    first f = lift $ dimap assoc1 assoc2 (first (unlift f)) 
    second f = lift $ dimap shuffle1 shuffle2 (second (unlift f))
    f &&& g = lmap (\x -> (x,x)) (f *** g)
    f *** g = lift $ dimap assoc1 shuffle3 (first (unlift f))
                     >>> rmap shuffle2 (second (unlift g))
    {-# INLINE arr #-}
    {-# INLINE first #-}
    {-# INLINE second #-}
    {-# INLINE (&&&) #-}
    {-# INLINE (***) #-}
  
  instance (ArrowChoice c, Profunctor c) => ArrowChoice (StateT s c) where
    left f = lift $ dimap distribute1 distribute2 (left (unlift f)) 
    right f = lift $ dimap distribute1 distribute2 (right (unlift f))
    f ||| g = lift $ lmap distribute1 (unlift f ||| unlift g)
    f +++ g = lift $ dimap distribute1 distribute2 (unlift f +++ unlift g)
    {-# INLINE left #-}
    {-# INLINE right #-}
    {-# INLINE (+++) #-}
    {-# INLINE (|||) #-}
  
  instance (ArrowApply c, Profunctor c) => ArrowApply (StateT s c) where
    app = lift $ lmap (\(s,(f,b)) -> (unlift f,(s,b))) app
    {-# INLINE app #-}
  
  instance (ArrowCont c) => ArrowCont (StateT s c) where
    type Cont (StateT s c) x = Cont c (s,x)
    callCC f = lift $ callCC $ \k -> unlift (f k)
    jump k = lift $ jump k
    {-# INLINE callCC #-}
    {-# INLINE jump #-}
  
  instance (Arrow c, Profunctor c) => ArrowState s (StateT s c) where
    get = lift (arr (\(a,()) -> (a,a)))
    put = lift (arr (\(_,s) -> (s,())))
    modify f = lift (dimap (\(s,x) -> (s,(x,s))) (\(_,(y,s)) -> (s,y)) (unlift f))
    {-# INLINE get #-}
    {-# INLINE put #-}
    {-# INLINE modify #-}
  
  instance ArrowReader r c => ArrowReader r (StateT s c) where
    ask = lift' Reader.ask
    local f = lift $ lmap shuffle1 (Reader.local (unlift f))
    {-# INLINE ask #-}
    {-# INLINE local #-}
  
  instance ArrowWriter w c => ArrowWriter w (StateT s c) where
    tell = lift' tell
    {-# INLINE tell #-}
  
  instance ArrowConst x c => ArrowConst x (StateT s c) where
    askConst f = lift (askConst (unlift . f))
    {-# INLINE askConst #-}
  
  second' :: (x -> y) -> ((z,x) -> (z,y))
  second' f (x,y) = (x,f y)
  {-# INLINE second' #-}
  