1 files changed, 102 insertions, 0 deletions

diff --git a/src/Data/PacketBuffer.hs b/src/Data/PacketBuffer.hs
new file mode 100644
index 00000000..c5ede50d
--- /dev/null
+++ b/src/Data/PacketBuffer.hs
@@ -0,0 +1,102 @@
+{-# LANGUAGE TupleSections #-}
+module Data.PacketBuffer
+    ( PacketBuffer
+    , newPacketBuffer
+    , PacketEvent(..)
+    , grokPacket
+    , awaitReadyPacket
+    , packetNumbersToRequest
+    , expectingSequenceNumber
+    , nextToSendSequenceNumber
+    , retrieveForResend
+    , decompressSequenceNumbers
+    ) where
+
+import Data.PacketQueue as Q
+
+import Control.Concurrent.STM
+import Control.Monad
+import Data.Maybe
+import Data.Word
+
+data PacketBuffer a b = PacketBuffer
+    { inQueue   :: PacketQueue a
+    , outBuffer :: PacketQueue b
+    }
+
+-- | Initialize the packet buffers.  Note, the capacity of the inbound packet
+-- queue is currently hardcoded to 200 packets and the capacity of the outbound
+-- buffer is hardcoded to 400 packets.
+newPacketBuffer :: STM (PacketBuffer a b)
+newPacketBuffer = PacketBuffer <$> Q.new 200 0
+                               <*> Q.new 400 0
+
+-- | Input for 'grokPacket'.
+data PacketEvent a b
+    = PacketSent { peSeqNum      :: Word32 -- ^ Sequence number for payload.
+                 , peSentPayload :: b      -- ^ Payload packet we sent to them.
+                 }
+    | PacketReceived { peSeqNum          :: Word32 -- ^ Sequence number for payload.
+                     , peReceivedPayload :: a      -- ^ Payload packet they sent to us.
+                     , peAcknowledgedNum :: Word32 -- ^ Earliest sequence number they've seen from us.
+                     }
+    | PacketReceivedLossy { peSeqNum          :: Word32 -- ^ Sequence number for lossy packet.
+                          , peAcknowledgedNum :: Word32 -- ^ Earliest sequence number they've seen from us.
+                          }
+
+-- | Whenever a packet is received or sent (but not resent) from the network,
+-- this function should be called to update the relevant buffers.
+--
+-- On outgoing packets, if the outbound buffer is full, this will block
+-- indefinitely until it is called in another thread with an inbound
+-- acknowledgement.
+grokPacket :: PacketBuffer a b -> PacketEvent a b -> STM ()
+grokPacket (PacketBuffer _   outb) (PacketSent seqno a)
+    = do (n,_) <- Q.enqueue outb seqno a
+         when (n/=0) retry
+grokPacket (PacketBuffer inb outb) (PacketReceived seqno a ack)
+    = do Q.enqueue inb seqno a
+         Q.dropPacketsBefore outb ack
+grokPacket (PacketBuffer inb outb) (PacketReceivedLossy seqno ack)
+    = do Q.observeOutOfBand inb seqno
+         Q.dropPacketsBefore outb ack
+
+-- | Wait until an inbound packet is ready to process.  Any necessary
+-- re-ordering will have been done.
+awaitReadyPacket :: PacketBuffer a b -> STM a
+awaitReadyPacket (PacketBuffer inb _) = Q.dequeue inb
+
+-- | Obtain a list of sequence numbers that may have been dropped.  This would
+-- be any number not yet received that is prior to the maxium sequence number
+-- we have received.  For convenience, a lowerbound for the missing squence numbers
+-- is also returned as the second item of the pair.
+packetNumbersToRequest :: PacketBuffer a b -> STM ([Word32],Word32)
+packetNumbersToRequest (PacketBuffer inb _) = do
+    ns <- Q.getMissing inb
+    lb <- Q.getLastDequeuedPlus1 inb
+    return (ns,lb)
+
+expectingSequenceNumber :: PacketBuffer a b -> STM Word32
+expectingSequenceNumber (PacketBuffer inb _ ) = Q.getLastDequeuedPlus1 inb
+
+nextToSendSequenceNumber :: PacketBuffer a b -> STM Word32
+nextToSendSequenceNumber (PacketBuffer _ outb) = Q.getLastEnqueuedPlus1 outb
+
+-- | Retrieve already-sent packets by their sequence numbers.  See
+-- 'decompressSequenceNumbers' to obtain the sequence number list from a
+-- compressed encoding.  There is no need to call 'grokPacket' when sending the
+-- packets returned from this call.
+retrieveForResend :: PacketBuffer a b -> [Word32] -> STM [(Word32,b)]
+retrieveForResend (PacketBuffer _ outb) seqnos =
+    catMaybes <$> forM seqnos (\no -> fmap (no,) <$> Q.lookup outb no)
+
+-- | Expand a compressed set of sequence numbers.  The first squence number is
+-- given directly and the rest are computed using 8-bit offsets.  This is
+-- normally used to obtain input for 'retrieveForResend'.
+decompressSequenceNumbers :: Word32 -> [Word8] -> [Word32]
+decompressSequenceNumbers baseno ns = foldr doOne (const []) ns (baseno-1)
+    where
+        doOne :: Word8 -> (Word32 -> [Word32]) -> Word32 -> [Word32]
+        doOne 0 f addend = f (addend+255)
+        doOne x f addend = let y = fromIntegral x + addend
+                           in y : f y