{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}

{- |
Module: Control.Concurrent.FVar
Description: Fast (partly lockless) mutable Vars
Copyright: (c) Erick Gonzalez, 2019
License: BSD3
Maintainer: erick@codemonkeylabs.de

This library implements an "FVar" or /fast Var/. The idea is that it allows concurrent readers
to access shared data without locking, but the moment a writer performs a write it gets exclusive
access to the var. If there are readers currently accesing the data, it waits until they complete
the operation, blocks any further reads and then it performs the write. Once it is done, it
reenables the stalled readers and lockless operation continues.

-}
module Control.Concurrent.FVar (modifyFVar, newFVar, readFVar, FVar) where

import Control.Concurrent.MVar
import Control.Monad                (replicateM_, when)
import Control.Monad.Catch          (bracket, bracket_, MonadMask)
import Control.Monad.IO.Class       (MonadIO, liftIO)
import Data.IORef
import Data.Maybe                   (isJust)
import System.IO.Unsafe             (unsafePerformIO)

-- | An 'FVar' (a.k.a. /eff-var/ or /Fast Var/) is a synchronising variable which allows
-- concurrent simultaneous read access without the need for locking. Write access is however
-- exclusive and causes all other operations to stall until the write is completed and
-- lockless operation can continue.
newtype FVar a = FVar { iRef :: IORef (RWVar a) } deriving (Eq)

data RWVar a = RWVar { readLock       :: Maybe (MVar ()),
                       readMutex      :: MVar (),
                       writeMutex     :: MVar (),
                       readersWaiting :: Int,
                       writersWaiting :: Int,
                       readers        :: Int,
                       writers        :: Int,
                       unboxed        :: a
                      }

#ifdef DEBUG
{-# NOINLINE out #-}
out :: MVar ()
out = unsafePerformIO $ newMVar ()
debug :: (MonadIO m) => String -> m ()
debug = liftIO . withMVar out . const . putStrLn
#else
debug :: (MonadIO m) => String -> m ()
debug = const $ return ()
#endif

-- | Creates a new 'FVar' containing the data provided. Note that unlike 'MVar's,
-- an 'FVar' can not be empty. This is to avoid complications stalling lock-less reads
-- in that case. There is of course nothing stopping you from storing a value wrapped in
-- a 'Maybe' to account for a similar usage pattern.
newFVar :: (MonadIO m) => a -> m (FVar a)
newFVar val = liftIO $ do
  rMVar <- newEmptyMVar
  wMVar <- newEmptyMVar
  ref  <- newIORef RWVar { readLock       = Nothing,
                           readersWaiting = 0,
                           writersWaiting = 0,
                           readers        = 0,
                           writers        = 0,
                           readMutex      = rMVar,
                           writeMutex     = wMVar,
                           unboxed        = val }
  return $ FVar ref

-- | Access an 'FVar' for reading using the provided function. Access is __not__ exclusive,
-- so multiple threads would be allowed to access the data simultaneously. Note that there is
-- actually nothing in principle preventing one from mutating shared data during this access
-- but in doing so you would be violating the contract in this API and thus integrity could
-- not be guaranteed, so you must refrain from doing so and truly only perform read accesses
-- to the shared data
readFVar :: (MonadMask m, MonadIO m) => FVar a -> (a -> m b) -> m b
readFVar FVar {..} = bracket wantsRead readDone
    where wantsRead = do
            (x, mLock) <- liftIO $ atomicModifyIORef' iRef readerIn
            maybe (access x) acquire mLock
          access x = do
            debug "reader reads fast"
            return x
          readDone _ = do
            debug "reader done"
            mLock <- liftIO $ atomicModifyIORef' iRef $ readerOut
            mapM_ release mLock
          acquire var  = do
            debug "reader waits"
            liftIO $ do
              takeMVar var
              debug "reader acquired lock"
              unboxed <$> readIORef iRef
          release var  = do
            debug "reader released write lock"
            liftIO $ putMVar var ()

-- | Modify the data inside an 'FVar'. Write access during execution of the function provided
-- will be exclusive so other writes or reads will be stalled until the call to 'modifyFVar'
-- is completed
modifyFVar :: (MonadMask m, MonadIO m) => FVar a -> (a -> m (a, b)) -> m b
modifyFVar FVar {..} mutator = bracket_ wantsWrite writeDone mutate
    where wantsWrite = liftIO $ mapM_ acquire =<< atomicModifyIORef' iRef writerIn
          writeDone  = liftIO $ do
            debug "writer done"
            mLock <- atomicModifyIORef' iRef writerOut
            mapM_ release mLock
          acquire var = do
            debug "writer waits"
            takeMVar var
            debug "writer acquired lock"
          release (waiters, var) = do
            debug "writer released lock"
            liftIO $ replicateM_ waiters $ putMVar var ()
          mutate = do
            debug "writer writes"
            (value, result) <- mutator =<< unboxed <$> liftIO (readIORef iRef)
            liftIO . atomicModifyIORef' iRef $ \var -> (var { unboxed = value }, result)

readerIn :: RWVar a -> (RWVar a, (a, Maybe (MVar ())))
readerIn var@RWVar {..} =
  (var { readers = readers + 1, readersWaiting = waitingNow }, (unboxed, readLock))
    where waitingNow | locked    = readersWaiting + 1
                     | otherwise = readersWaiting
          locked = isJust readLock

readerOut :: RWVar a -> (RWVar a, Maybe (MVar ()))
readerOut var@RWVar {..} =
  (var { readers = readersNow, writersWaiting = waitingNow}, withLock)
    where withLock
            | unlockWriter = return writeMutex
            | otherwise    = Nothing
          readersNow       = readers - 1
          waitingNow
            | unlockWriter = writersWaiting - 1
            | otherwise    = writersWaiting
          unlockWriter = readersNow == readersWaiting && writersWaiting /= 0

writerIn :: RWVar a -> (RWVar a, Maybe (MVar ()))
writerIn var@RWVar {..} =
  (var { readLock       = return readMutex,
         writers        = writers + 1,
         writersWaiting = waitingNow }, withLock)
    where waitingNow
            | stalled   = writersWaiting + 1
            | otherwise = writersWaiting
          withLock
            | stalled   = return writeMutex
            | otherwise = Nothing
          stalled = writers /= 0 || readers /= readersWaiting

writerOut :: RWVar a -> (RWVar a, Maybe (Int, MVar ()))
writerOut var@RWVar {..} =
  (var { readLock       = rearmLock,
         readersWaiting = 0,
         writers        = writers - 1,
         writersWaiting = nowWaiting }, withLock)
    where rearmLock
            | stalledWriters = return readMutex
            | otherwise      = Nothing
          withLock
            | stalledReaders = return (readersWaiting, readMutex)
            | stalledWriters = return (1, writeMutex)
            | otherwise      = Nothing
          nowWaiting
            | stalledWriters && not stalledReaders = writersWaiting - 1
            | otherwise                            = writersWaiting
          stalledReaders = readersWaiting /= 0
          stalledWriters = writersWaiting /= 0