RingBuffer.hs

-- | This is a thread-safe implementation of a mutable ring-buffer
-- built upon @vector@.

module Data.RingBuffer ( RingBuffer
                       , new
                       , clear
                       , append
                       , concat
                       , capacity
                       , length
                       , latest
                       , toList
                       , withItems
                       ) where

import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable as VGM
import Control.Applicative
import Control.Concurrent
import Control.Monad (when)
import Control.Monad.Catch
import Control.Monad.State
import Control.Monad.Reader
import Control.Monad.Primitive
import Prelude hiding (length, concat)

-- | A concurrent ring buffer.
data RingBuffer v a
    = RingBuffer { ringBuffer :: !(VG.Mutable v (PrimState IO) a)
                 , ringState  :: !(MVar RingState)
                 }

data RingState = RingState { ringFull :: !Bool -- ^ is the ring full?
                           , ringHead :: !Int -- ^ index of next entry to be written
                           }

-- | We use the @Mutable@ vector type to ensure injectiveness
type RingM m vm a = StateT RingState (ReaderT (vm (PrimState IO) a) m)

-- | Atomically perform an action with the ring
withRing :: (VG.Vector v a, MonadIO m, MonadMask m)
         => RingBuffer v a
         -> RingM m (VG.Mutable v) a r
         -> m r
withRing rb action = mask_ $ do
    s <- liftIO $ takeMVar (ringState rb)
    (r, s') <- runReaderT (runStateT action s) (ringBuffer rb)
    liftIO $ putMVar (ringState rb) s'
    return r

advance :: (VGM.MVector v a, MonadIO m) => Int -> RingM m v a ()
advance n = do
    RingState full pos <- get
    cap <- capacity'
    let (a, pos') = (pos + n) `divMod` cap
    put $ RingState (full || a > 0) pos'
{-# INLINEABLE advance #-}

-- | Create a new ring of a given length
--
-- /Note:/ size must be non-zero
new :: (VG.Vector v a) => Int -> IO (RingBuffer v a)
new n = do
    when (n < 1) $ fail "Data.RingBuffer.new: invalid ring size"
    buffer <- VGM.new n
    state0 <- newMVar $ RingState False 0
    return $ RingBuffer { ringBuffer=buffer, ringState=state0 }

-- | Reset the ringbuffer to its empty state
clear :: VG.Vector v a => RingBuffer v a -> IO ()
clear rb = withRing rb $ put $ RingState False 0

-- | Add an item to the end of the ring
append :: (VG.Vector v a) => a -> RingBuffer v a -> IO ()
append x rb = withRing rb $ do
    s <- get
    liftIO $ VGM.unsafeWrite (ringBuffer rb) (ringHead s) x
    advance 1
{-# INLINABLE append #-}

-- | Add multiple items to the end of the ring
-- This ignores any items above the length of the ring
concat :: (VG.Vector v a) => v a -> RingBuffer v a -> IO ()
concat xs rb = withRing rb $ do
    cap <- capacity'
    let takeN = min (VG.length xs) cap
    xs' <- liftIO $ VG.unsafeThaw $ VG.drop (VG.length xs - takeN) xs
    pos <- gets ringHead

    let untilWrap = cap - pos
        src  = VGM.take untilWrap xs'
        dest = VGM.take (min takeN untilWrap) $ VGM.drop pos $ ringBuffer rb
    liftIO $ VGM.copy dest src

    -- did we wrap around?
    when (takeN > untilWrap) $ do
        let src'  = VGM.drop untilWrap xs'
            dest' = VGM.take (takeN - untilWrap) $ ringBuffer rb
        liftIO $ VGM.copy dest' src'
    advance takeN
{-# INLINABLE concat #-}

-- | The maximum number of items the ring can contain
capacity :: (VG.Vector v a) => RingBuffer v a -> Int
capacity rb = VGM.length (ringBuffer rb)

-- | The maximum number of items the ring can contain
capacity' :: (VGM.MVector v a, MonadIO m) => RingM m v a Int
capacity' = asks VGM.length

-- | The current filled length of the ring
length' :: (VGM.MVector v a, MonadIO m) => RingM m v a Int
length' = do
    RingState full pos <- get
    if full
      then capacity'
      else return pos

-- | The current filled length of the ring
length :: (VG.Vector v a) => RingBuffer v a -> IO Int
length rb = withRing rb length'

-- | Retrieve the \(n\)th most-recently added item of the ring
latest :: (VG.Vector v a) => RingBuffer v a -> Int -> IO (Maybe a)
latest rb n = withRing rb $ do
    len <- length'
    if n >= len
      then return Nothing
      else Just <$> latest' n

latest' :: (VGM.MVector v a, MonadIO m) => Int -> RingM m v a a
latest' n = do
    len <- length'
    cap <- capacity'
    when (n >= len) $ error "Data.RingBuffer.latest': invalid index"
    RingState _ hd <- get
    let idx = (hd - n - 1) `mod` cap
    buf <- ask
    liftIO $ VGM.unsafeRead buf idx
{-# INLINABLE latest' #-}

-- | Get the entire contents of the ring, with the most recently added element
-- at the head. Note that this is rather inefficient.
toList :: (VG.Vector v a) => RingBuffer v a -> IO [a]
toList rb = withRing rb $ do
    len <- length'
    mapM latest' [0..len-1]
{-# INLINABLE toList #-}

-- | Execute the given action with the items of the ring.
-- Note that no references to the vector may leak out of the action as
-- it will later be mutated. Moreover, the items in the vector are in
-- no particular order.
withItems :: (MonadIO m, MonadMask m, VG.Vector v a)
          => RingBuffer v a -> (v a -> m b) -> m b
withItems rb action = withRing rb $ do
    frozen <- liftIO $ VG.unsafeFreeze (ringBuffer rb)
    n <- length'
    lift $ lift $ action (VG.take n frozen)