module Data.Queue
    ( Queue
    , Data.Queue.empty
    , Input
    , Output
    , isEmpty
    , snoc
    , pop
    , cons
    , append
    , Data.Queue.foldr
    , Data.Queue.map
    , send
    , recv
    , sendAtomic
    , recvAtomic) where

import qualified Data.DList as D
import Control.Concurrent.STM

-- |A buffer data type with the direction ('Input' or 'Output') and element type.
data Queue dir a = Q (D.DList a) Int

data Input
data Output

empty :: Queue d a
empty = Q D.empty 0

-- |Tests if the element count of a queue is zero.
isEmpty :: Queue d a -> Bool
isEmpty (Q _ 0) = True
isEmpty _ = False

-- |Places an element at the end of the queue. O(1)
snoc :: Queue d a -> a -> Queue d a
snoc (Q d n) e = Q (D.snoc d e) (n + 1)

-- |Extracts an element from the front of the queue. O(1)
pop :: Queue d a -> (Maybe a, Queue d a)
pop (Q d 0) = (Nothing, Q d 0)
pop (Q d n) = (Just (D.head d), Q (D.tail d) (n - 1))

-- |Places an element at the front of the queue. O(1)
cons :: a -> Queue d a -> Queue d a
cons e (Q d n) = Q (D.cons e d) (n + 1)

-- |Combine two queues.
append :: Queue d a -> Queue d a -> Queue d a
append (Q q i) (Q p j) = Q (D.append q p) (i + j)

map :: (a -> b) -> Queue d a -> Queue d b
map f (Q xs len) = Q (D.map f xs) len

foldr :: (a -> b -> b) -> b -> Queue d a -> b
foldr f z (Q xs len) = D.foldr f z xs

-- |IO action for snoc on a TVar queue
send :: TVar (Queue Output a) -> a -> IO ()
send tv a = atomically ( sendAtomic tv a )

-- |STM action to snoc on a TVar queue
sendAtomic :: TVar (Queue Output a) -> a -> STM ()
sendAtomic tv a = do
    q <- readTVar tv
    writeTVar tv (snoc q a)

-- |IO action for popping an element from a TVar queue
recv :: TVar (Queue Input a) -> IO a
recv tv = atomically (recvAtomic tv)

-- |STM action for popping an element from a TVar queue
recvAtomic :: TVar (Queue Input a) -> STM a
recvAtomic tv = do
    q <- readTVar tv
    case pop q of
        (Nothing, _) -> retry
        (Just e, qs) -> writeTVar tv qs >> return e