-- | This module is useful for implementing a lossess protocol on top of a
-- lossy datagram style protocol.  It implements a buffer in which packets may
-- be stored out of order, but from which they are extracted in the proper
-- sequence.
{-# LANGUAGE NamedFieldPuns #-}
module Data.PacketQueue
    ( PacketQueue
    , new
    , dequeue
    , enqueue
    , observeOutOfBand
    ) where

import Control.Concurrent.STM
import Control.Concurrent.STM.TArray
import Control.Monad
import Data.Word
import Data.Array.MArray

data PacketQueue a = PacketQueue
    { pktq    :: TArray Word32 (Maybe a)
    , seqno   :: TVar Word32
    , qsize   :: Word32
    , buffend :: TVar Word32
    }

-- | Create a new PacketQueue.
new :: Word32 -- ^ Capacity of queue.
    -> Word32 -- ^ Initial sequence number.
    -> STM (PacketQueue a)
new capacity seqstart = do
    let cap = if capacity `mod` 2 == 0 then capacity else capacity + 1
    q <- newArray (0,cap - 1) Nothing
    seqv <- newTVar seqstart
    bufe <- newTVar 0
    return PacketQueue
        { pktq    = q
        , seqno   = seqv
        , qsize   = cap
        , buffend = bufe
        }

observeOutOfBand :: PacketQueue a -> Word32-> STM ()
observeOutOfBand PacketQueue { seqno, qsize, buffend } no = do
    low <- readTVar seqno
    let proj = no - low
    -- Ignore packet if out of range.
    when ( proj < qsize) $ do
    modifyTVar' buffend (\be -> if be - low <= proj then no + 1 else be)


-- | Retry until the next expected packet is enqueued.  Then return it.
dequeue :: PacketQueue a -> STM a
dequeue PacketQueue { pktq, seqno, qsize } = do
    i0 <- readTVar seqno
    let i = i0 `mod` qsize
    x <- maybe retry return =<< readArray pktq i
    writeArray pktq i Nothing
    modifyTVar' seqno   succ
    return x

-- | Enqueue a packet.  Packets need not be enqueued in order as long as there
-- is spare capacity in the queue.  If there is not, the packet will be
-- silently discarded without blocking.
enqueue :: PacketQueue a -- ^ The packet queue.
        -> Word32        -- ^ Sequence number of the packet.
        -> a             -- ^ The packet.
        -> STM ()
enqueue PacketQueue{ pktq, seqno, qsize, buffend } no x = do
    low <- readTVar seqno
    let proj = no - low
    -- Ignore packet if out of range.
    when ( proj < qsize) $ do
    let i = no `mod` qsize
    writeArray pktq i (Just x)
    modifyTVar' buffend (\be -> if be - low <= proj then no + 1 else be)

-- lookup :: PacketQueue a -> Word32 -> STM (Maybe a)
-- lookup PacketQueue{ pktq, seqno, qsize } no = _todo