{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE FunctionalDependencies     #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE RankNTypes                 #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Control.Monad.Queue.Corec
-- Copyright   :  (c) Leon P Smith 2009
-- License     :  BSD3
--
-- Maintainer  :  leon at melding-monads dot com
-- Stability   :  experimental
-- Portability :  portable
--
-- Corecursive queues with return values.  This is a straightforward
-- generalization of Control.Monad.Queue.Allison.   It corresponds to
-- @CorecQW@ in the paper
-- /Lloyd Allison's Corecursive Queues:  Why Continuations Matter/ by
-- Leon P Smith in the Monad Reader issue 14.
--
-----------------------------------------------------------------------------

module  Control.Monad.Queue.Corec
     (  Q()
     ,  LenType
     ,  runResultQueue
     ,  runResult
     ,  runQueue
     ,  enQ
     ,  deQ
     ,  deQ_break
     ,  deQs
     ,  peekQ
     ,  peekQn
     ,  peekQs
     ,  lenQ
     ,  lenQ_
     ,  mapQ
     ,  wfix
     ,  callCC
     ,  exit
     )  where

import qualified Control.Monad.Queue.Class as Class
import Control.Monad.Queue.Util
import Data.List(genericIndex, genericTake, genericSplitAt)

type QSt w e = LenType -> [e] -> (w,[e])

newtype Q w e a = Q { unQ :: (a -> QSt w e) -> QSt w e }

instance Monad (Q w e) where
  return a  = Q (\k -> k a)
  m >>= f   = Q (\k -> unQ m (\a -> unQ (f a) k))

callCC :: ((a -> forall b. Q w e b) -> Q w e a) -> Q w e a
callCC f = Q $ \c -> unQ (f (\a -> Q $ \_ -> c a)) c

-- | Terminates the queue computation with result @w@
exit :: w -> Q w e a
exit w = Q (\_k _n _q -> (w,[]))


-- | Enqueues an element to the queue
enQ :: e -> Q w e ()
enQ e = Q (\k n q -> let (w,es) = (k () $! n+1) q
                      in (w,e:es))

-- | Dequeues and element:  returns 'Nothing' if the queue is empty.
deQ  ::  Q w e (Maybe e)
deQ  =   Q delta
  where
    delta k n q
       | n <= 0     = k Nothing n q
       | otherwise  = case q of
                       [] -> error "Control.Monad.Queue.Corec.deQ: empty list"
                       (e:q') -> (k (Just e) $! n-1) q'

-- | Dequeues an element:  terminates the computation with the final result @w@ if the queue is empty.
deQ_break :: w -> Q w e e
deQ_break w =   Q delta
  where
    delta k n q
       | n <= 0    = (w,[])
       | otherwise = case q of
                      [] -> error "Control.Monad.Queue.Corec.deQ_break: empty list"
                      (e:q') -> (k e $! n-1) q'

-- | Dequeues up to @len@ elements from the queue
deQs :: Integral len => len -> Q w e [e]
deQs i = Q delta
  where
    delta k n q
       = let i' = min (fromIntegral (max 0 i)) n
             (res,q') = genericSplitAt i' q
          in (k res $! n-i') q'

-- | Examines the front element of the queue without removing it.
peekQ :: Q w e (Maybe e)
peekQ = Q delta
  where
    delta k n q
       | n <= 0    = k Nothing n q
       | otherwise = case q of
                      [] -> error "Control.Monad.Queue.Corec.peekQ: empty list"
                      (e:q') -> k (Just e) n q

-- | Examines the element currently at position @index@ in the queue,  indexing starts from @0@, like '!!'
peekQn :: (Integral index) => index -> Q w e (Maybe e)
peekQn i_ = Q delta
  where
    i = fromIntegral i_

    delta k n q
       | not (0 <= i_ && i < n) = k Nothing n q
       | otherwise = k (Just (genericIndex q i)) n q

-- | Looks at up to the first @len@ elements of the queue,  like 'deQs' except without removing them.
peekQs :: (Integral len) => len -> Q w e [e]
peekQs i_ = Q delta
  where
    i = fromIntegral (max 0 i_)
    delta k n q = k (genericTake (min i n) q) n q

-- | Returns the length of the queue
lenQ_ :: Q w e LenType
lenQ_ = Q (\k n q -> k n n q)

-- | Returns the length of the queue
lenQ  :: Integral len => Q w e len
lenQ  = Q (\k n q -> k (fromIntegral n) n q)

-- | Applies a function to the final return value of the entire computation,  like 'Control.Monad.Cont.mapCont'
mapQ :: (w -> w) -> Q w e a -> Q w e a
mapQ f m = Q (\k n q ->  let (w,es) = unQ m k n q
                         in  (f w, es))

-- | Computes a fixpoint on the result;  usually used in conjunction with @mapQ@
wfix :: (w -> Q w e a) -> Q w e a
wfix f = Q (\k n q -> let (w,q') = unQ (f w) k n q
                       in (w,q'))


-- | Runs the computation,  returns the result of the computation and a list of all elements enqueued
runResultQueue :: Q a e a -> (a,[e])
runResultQueue m = st
  where
    st@(_a,q) = unQ m (\a _ _ -> (a,[])) 0 q


-- | Runs the computation,  returns the result of the computation
--
runResult :: Q a e a -> a
runResult = fst . runResultQueue

-- | Runs the computation,  returns a list of all elements enqueued
runQueue :: Q a e a -> [e]
runQueue  = snd . runResultQueue

instance Class.MonadQueue e (Q w e) where
  enQ     = enQ
  peekQ   = peekQ
  peekQs  = peekQs
  peekQn  = peekQn
  deQ     = deQ
  deQs    = deQs
  lenQ    = lenQ