1 files changed, 44 insertions, 7 deletions

diff --git a/src/Data/PacketQueue.hs b/src/Data/PacketQueue.hs
index a4c99cab..4f0f04e3 100644
--- a/src/Data/PacketQueue.hs
+++ b/src/Data/PacketQueue.hs
@@ -6,8 +6,11 @@
 {-# LANGUAGE FlexibleContexts #-}
 module Data.PacketQueue
     ( PacketQueue
+    , getCapacity
     , new
+    , newOverwrite
     , dequeue
+    , markButNotDequeue
     , enqueue
     , observeOutOfBand
     , PacketOutQueue
@@ -24,6 +27,7 @@ module Data.PacketQueue
 import Control.Concurrent.STM
 import Control.Concurrent.STM.TArray
 import Control.Monad
+import Control.Applicative
 import Data.Word
 import Data.Array.MArray
 import Data.Maybe
@@ -33,6 +37,7 @@ data PacketQueue a = PacketQueue
     , seqno   :: TVar Word32
     , qsize   :: Word32
     , buffend :: TVar Word32 -- on incoming, highest packet number handled + 1
+    , qOverWriteMode :: Bool
     }
 
 packetQueueViewList :: PacketQueue a -> STM [(Word32,a)]
@@ -41,6 +46,8 @@ packetQueueViewList p = do
         f (n,Just x)  = Just (n,x)
     catMaybes . map f <$> getAssocs (pktq p)
 
+getCapacity :: Applicative m => PacketQueue t -> m Word32
+getCapacity (PacketQueue { qsize }) = pure qsize
 
 -- | Create a new PacketQueue.
 new :: Word32 -- ^ Capacity of queue.
@@ -56,6 +63,24 @@ new capacity seqstart = do
         , seqno   = seqv
         , qsize   = cap
         , buffend = bufe
+        , qOverWriteMode = False
+        }
+
+-- | Create a new PacketQueue with Overwrite on Wrap.
+newOverwrite :: Word32 -- ^ Capacity of queue.
+    -> Word32 -- ^ Initial sequence number.
+    -> STM (PacketQueue a)
+newOverwrite 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
+        , qOverWriteMode = True
         }
 
 observeOutOfBand :: PacketQueue a -> Word32-> STM ()
@@ -77,21 +102,33 @@ dequeue PacketQueue { pktq, seqno, qsize } = do
     modifyTVar' seqno   succ
     return x
 
+-- | Like dequeue, but marks as handled rather than removing
+markButNotDequeue :: PacketQueue (Bool,a) -> STM a
+markButNotDequeue PacketQueue { pktq, seqno, qsize } = do
+    i0 <- readTVar seqno
+    let i = i0 `mod` qsize
+    (b,x) <- maybe retry return =<< readArray pktq i
+    writeArray pktq i (Just (True,x))
+    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.
+-- silently discarded without blocking. (Unless this is an Overwrite-queue,
+-- in which case, the packets will simply wrap around overwriting the old ones.)
 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
+        -> STM (Word32,Word32)
+enqueue PacketQueue{ pktq, seqno, qsize, buffend,qOverWriteMode } 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)
+    when ( proj < qsize || qOverWriteMode) $ do
+        let i = no `mod` qsize
+        writeArray pktq i (Just x)
+        modifyTVar' buffend (\be -> if be - low <= proj then no + 1 else be)
+    return (proj `divMod` qsize)
 
 -- lookup :: PacketQueue a -> Word32 -> STM (Maybe a)
 -- lookup PacketQueue{ pktq, seqno, qsize } no = _todo