{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances #-}

-- ------------------------------------------------------------

{- |
   Module     : Control.Arrow.IOStateListArrow
   Copyright  : Copyright (C) 2005-8 Uwe Schmidt
   License    : MIT

   Maintainer : Uwe Schmidt (uwe\@fh-wedel.de)
   Stability  : experimental
   Portability: portable

   Implementation of arrows with IO and a state

-}

-- ------------------------------------------------------------

module Control.Arrow.IOStateListArrow
    ( IOSLA(..)
    , liftSt
    , runSt
    )
where

import           Prelude hiding (id, (.))

import           Control.Category

import           Control.Arrow
import           Control.Arrow.ArrowIf
import           Control.Arrow.ArrowIO
import           Control.Arrow.ArrowList
import           Control.Arrow.ArrowNF
import           Control.Arrow.ArrowTree
import           Control.Arrow.ArrowState

import           Control.DeepSeq

-- ------------------------------------------------------------

-- | list arrow combined with a state and the IO monad

newtype IOSLA s a b = IOSLA { runIOSLA :: s -> a -> IO (s, [b]) }

instance Category (IOSLA s) where
    id                  = IOSLA $ \ s x -> return (s, [x])	-- don't defined id = arr id, this gives loops during optimization

    IOSLA g . IOSLA f	= IOSLA $ \ s x -> do
					   (s1, ys) <- f s x
					   sequence' s1 ys
					   where
					   sequence' s' []	 = return (s', [])
					   sequence' s' (x':xs') = do
								   (s1', ys') <- g s' x'
								   (s2', zs') <- sequence' s1' xs'
								   return (s2', ys' ++ zs')

instance Arrow (IOSLA s) where
    arr f		= IOSLA $ \ s x -> return (s, [f x])

    first (IOSLA f)	= IOSLA $ \ s (x1, x2) -> do
					           (s', ys1) <- f s x1
					           return (s', [ (y1, x2) | y1 <- ys1 ])

    -- just for efficiency
    second (IOSLA g)	= IOSLA $ \ s (x1, x2) -> do
					          (s', ys2) <- g s x2
					          return (s', [ (x1, y2) | y2 <- ys2 ])

    -- just for efficiency
    IOSLA f *** IOSLA g	= IOSLA $ \ s (x1, x2) -> do
						   (s1, ys1) <- f s  x1
					           (s2, ys2) <- g s1 x2
					           return (s2, [ (y1, y2) | y1 <- ys1, y2 <- ys2 ])

    -- just for efficiency
    IOSLA f &&& IOSLA g	= IOSLA $ \ s x -> do
					   (s1, ys1) <- f s  x
					   (s2, ys2) <- g s1 x
					   return (s2, [ (y1, y2) | y1 <- ys1, y2 <- ys2 ])



instance ArrowZero (IOSLA s) where
    zeroArrow		= IOSLA $ \ s -> const (return (s, []))


instance ArrowPlus (IOSLA s) where
    IOSLA f <+> IOSLA g	= IOSLA $ \ s x -> do
					   (s1, rs1) <- f s  x
					   (s2, rs2) <- g s1 x
					   return (s2, rs1 ++ rs2)

instance ArrowChoice (IOSLA s) where
    left (IOSLA f)	= IOSLA $ \ s -> either
			                 (\ x -> do
					         (s1, y) <- f s x
					         return (s1, map Left y)
					 )
                                         (\ x -> return (s, [Right x]))

    right (IOSLA f)	= IOSLA $ \ s -> either
                                         (\ x -> return (s, [Left x]))
			                 (\ x -> do
					         (s1, y) <- f s x
					         return (s1, map Right y)
					 )

instance ArrowApply (IOSLA s) where
    app			= IOSLA $ \ s (IOSLA f, x) -> f s x

instance ArrowList (IOSLA s) where
    arrL f		= IOSLA $ \ s x -> return (s, (f x))
    arr2A f		= IOSLA $ \ s (x, y) -> runIOSLA (f x) s y
    constA c		= IOSLA $ \ s   -> const (return (s, [c]))
    isA p		= IOSLA $ \ s x -> return (s, if p x then [x] else [])
    IOSLA f >>. g	= IOSLA $ \ s x -> do
				           (s1, ys) <- f s x
					   return (s1, g ys)
    -- just for efficency
    perform (IOSLA f)	= IOSLA $ \ s x -> do
					   (s1, _ys) <- f s x
					   return (s1, [x])

instance ArrowIf (IOSLA s) where
    ifA (IOSLA p) ta ea	= IOSLA $ \ s x -> do
				           (s1, res) <- p s x
					   runIOSLA ( if null res
						      then ea
						      else ta
						    ) s1 x

    (IOSLA f) `orElse` g
			= IOSLA $ \ s x -> do
				           r@(s1, res) <- f s x
					   if null res
					      then runIOSLA g s1 x
					      else return r


instance ArrowIO (IOSLA s) where
    arrIO cmd		= IOSLA $ \ s x -> do
					   res <- cmd x
					   return (s, [res])

instance ArrowIOIf (IOSLA s) where
    isIOA p		= IOSLA $ \ s x -> do
					   res <- p x
					   return (s, if res then [x] else [])

instance ArrowState s (IOSLA s) where
    changeState cf	= IOSLA $ \ s x -> let s' = cf s x in return (seq s' s', [x])
    accessState	af	= IOSLA $ \ s x -> return (s, [af s x])

-- ------------------------------------------------------------

-- |
-- lift the state of an IOSLA arrow to a state with an additional component.
--
-- This is uesful, when running predefined IO arrows, e.g. for document input,
-- in a context with a more complex state component.

liftSt		:: IOSLA s1 b c -> IOSLA (s1, s2) b c
liftSt (IOSLA f)
    = IOSLA $ \ (s1, s2) x -> do
			      (s1', ys) <- f s1 x
			      return ((s1', s2), ys)


-- |
-- run an arrow with augmented state in the context of a simple state arrow.
-- An initial value for the new state component is needed.
--
-- This is useful, when running an arrow with an extra environment component, e.g.
-- for namespace handling in XML.

runSt		:: s2 -> IOSLA (s1, s2) b c -> IOSLA s1 b c
runSt s2 (IOSLA f)
    = IOSLA $ \ s1 x -> do
			((s1', _s2'), ys) <- f (s1, s2) x
			return (s1', ys)

-- ------------------------------------------------------------

instance ArrowTree (IOSLA s)

instance (NFData s) => ArrowNF (IOSLA s) where
    rnfA (IOSLA f)	= IOSLA $ \ s x -> do
					   res <- f s x
					   deepseq res $ return res

-- ------------------------------------------------------------