MSQueue.hs

{-# LANGUAGE BangPatterns, MagicHash #-}

{-|
Module      : MSQueue
Description : Michael-Scott queue.
License     : BSD3
Maintainer  : Julian Sutherland (julian.sutherland10@imperial.ac.uk)

An Implementation of a Michael-Scott Lock-Free queues.
-}

module Data.NonBlocking.LockFree.MSQueue(MSQueue(), MSQueueIO, MSQueueSTM, newMSQueue, dequeueMSQueue, enqueueMSQueue) where

import Control.Concurrent.STM (STM())
import Control.Concurrent.STM.TVar (TVar())
import Control.Monad (when)
import Control.Monad.Loops (whileM_)
import Control.Monad.Ref (MonadAtomicRef, newRef, readRef, writeRef, atomicModifyRef)
import Data.IORef (IORef())
import Data.Maybe(isNothing, fromJust)
import GHC.Exts (Int(I#))
import GHC.Prim (reallyUnsafePtrEquality#)

-- |'MSQueue' inside the 'IO' 'Monad'.
type MSQueueIO a = MSQueue IORef a
-- |'MSQueue' inside the 'STM' 'Monad'.
type MSQueueSTM a = MSQueue TVar a

-- |Implementation of Michael-Scott Lock-Free queues. Specification and pseudo-code can be found at <http://www.research.ibm.com/people/m/michael/podc-1996.pdf>. Works with any combination of 'Monad' and reference satsfying the "MonadAtomicRef" class.
data MSQueue r a = MSQueue (r (r (Maybe (MSQueueElem r a)))) (r (r (Maybe (MSQueueElem r a))))
data MSQueueElem r a = MSQueueElem a (r (Maybe (MSQueueElem r a)))

instance Eq a => Eq (MSQueueElem r a) where
  (MSQueueElem x r1) == (MSQueueElem y r2) = (x == y) && ptrEq r1 r2

-- |Creates a new instance of 'MSQueue'. Internally implemented with a reference of type r, which is why they must be atomically modifiable. Initially empty.
{-# SPECIALIZE newMSQueue :: IO (MSQueueIO a) #-}
{-# SPECIALIZE newMSQueue :: STM (MSQueueSTM a) #-}
newMSQueue :: (MonadAtomicRef r m) => m (MSQueue r a)
newMSQueue = do
  null <- newRef Nothing
  topRefRef <- newRef null
  lastRefRef <- newRef null
  return (MSQueue topRefRef lastRefRef)

-- |Dequeues an element from a 'MSQueue' in a lock-free manner. Returns Nothing if the queue is empty, otherwise return the element wrapped in a 'Just'.
{-# SPECIALIZE dequeueMSQueue :: MSQueueIO a -> IO (Maybe a) #-}
{-# SPECIALIZE dequeueMSQueue :: MSQueueSTM a -> STM (Maybe a) #-}
dequeueMSQueue :: (MonadAtomicRef r m) => MSQueue r a -> m (Maybe a)
dequeueMSQueue (MSQueue topRefRef _) = do
  res <- newRef Nothing
  whileM_ (readRef res >>= return . isNothing) $ do
    topRef <- readRef topRefRef
    top    <- readRef topRef
    case top of
      (Just (MSQueueElem v tailRef)) -> do
        b <- casRef topRefRef topRef tailRef
        when b $ writeRef res (Just (Just v))
      Nothing -> writeRef res (Just Nothing)
  readRef res >>= return . fromJust

-- |Enqueues an element in a 'MSQueue' in a lock-free manner.
{-# SPECIALIZE enqueueMSQueue :: MSQueueIO a -> a -> IO () #-}
{-# SPECIALIZE enqueueMSQueue :: MSQueueSTM a -> a -> STM () #-}
enqueueMSQueue :: (MonadAtomicRef r m) => MSQueue r a -> a -> m ()
enqueueMSQueue (MSQueue _ lastRefRef) v = do
  nLastRef <- newRef Nothing
  let nLastElem = MSQueueElem v nLastRef
  done <- newRef False
  whileM_ (readRef done >>= return . not) $ do
    lastRef <- readRef lastRefRef
    last <- readRef lastRef
    if (isNothing last)
      then do
        b <- atomicModifyRef lastRef (\val -> let b = isNothing val in (if b then (Just nLastElem) else val, b))
        when b $ do
          casRef lastRefRef lastRef nLastRef
          writeRef done True
      else do
        let (Just (MSQueueElem _ nextRef)) = last
        casRef lastRefRef lastRef nextRef
        return ()

{-# SPECIALIZE casRef :: IORef (IORef a) -> IORef a -> IORef a -> IO Bool #-}
{-# SPECIALIZE casRef :: TVar (TVar a) -> TVar a -> TVar a -> STM Bool #-}
casRef :: (MonadAtomicRef r m) => r (r a) -> r a -> r a -> m Bool
casRef ref comp rep = atomicModifyRef ref (\val -> let b = ptrEq val comp in (if b then rep else val, b))

{-# NOINLINE ptrEq #-}
ptrEq :: a -> a -> Bool
ptrEq !x !y = I# (reallyUnsafePtrEquality# x y) == 1