content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  module Util.WBVar (
    WBVar,
    newWBVar, newWBVarIO,
    readWBVar, writeWBVar, modifyWBVar', dirtyWBVar,
    writeBackWBVar,
    WBMVar,
    newEmptyWBMVar,
    newWBMVar,
    readWBMVar,
    takeWBMVar,
    writeWBMVar,
    withWBMVar,
    modifyWBMVar,
  ) where
  
  import Util.Defer
  
  import Control.Concurrent (forkIO)
  import Control.Concurrent.STM
  import Control.Exception
  
  data WBVar a = WBVar
    { wbValue :: TVar a
    , wbDirty :: TVar Bool
    , wbWriting :: TVar Bool
    , wbWriteBack :: a -> IO ()
    }
  
  newWBVar :: a -> (a -> IO ()) -> STM (WBVar a)
  newWBVar x wb = WBVar
    <$> newTVar x
    <*> newTVar False
    <*> newTVar False
    <*> pure wb
  
  newWBVarIO :: a -> (a -> IO ()) -> IO (WBVar a)
  newWBVarIO x wb = WBVar
    <$> newTVarIO x
    <*> newTVarIO False
    <*> newTVarIO False
    <*> pure wb
  
  readWBVar :: WBVar a -> STM a
  readWBVar = readTVar . wbValue
  
  writeWBVar :: WBVar a -> a -> Defer IO STM ()
  writeWBVar v x = do
    lift $ writeTVar (wbValue v) x
    dirtyWBVar v
  
  modifyWBVar' :: WBVar a -> (a -> a) -> Defer IO STM ()
  modifyWBVar' v f = do
    lift $ modifyTVar' (wbValue v) f
    dirtyWBVar v
  
  dirtyWBVar :: WBVar a -> Defer IO STM ()
  dirtyWBVar v = do
    lift $ writeTVar (wbDirty v) True
    later $ writeBackWBVar v
  
  -- Write back the variable if it is dirty.
  --
  -- NOTE: To avoid blocking an arbitrary number of threads, we'll keep writing
  -- until the variable is no longer dirty. Suppose thread T1 is writing back
  -- and T2 dirties the variable while this is happening and tries to write back
  -- as well. In this situation, T2 will immediately return instead of blocking
  -- until T1 is done. T1, on the other hand, will finish writing back, notice
  -- that the variable is dirty and write back again.
  writeBackWBVar :: WBVar a -> IO ()
  writeBackWBVar v = do
    r <- atomically $ do
      dirty <- readTVar $ wbDirty v
      if dirty
        then do
          writing <- readTVar $ wbWriting v
          if writing
            then return Nothing
            else do
              writeTVar (wbDirty v) False
              writeTVar (wbWriting v) True
              Just <$> readTVar (wbValue v)
        else return Nothing
    case r of
      Just x -> do
        wbWriteBack v x
          `finally` atomically (writeTVar (wbWriting v) False)
          -- FIXME: A lame attempt to keep writing even in the case of an
          -- exception. I don't think this can really work without atomicity
          -- between the STM transaction and the write-back which seems impossible
          -- to achieve with STM.
          `onException` forkIO (writeBackWBVar v)
        writeBackWBVar v
      Nothing -> return ()
  
  
  -- | A 'WBVar' that can be empty/full like 'MVar'
  type WBMVar a = WBVar (Maybe a)
  
  newWBMVar :: a -> (a -> IO ()) -> STM (WBMVar a)
  newWBMVar a f = newWBVar (Just a) (maybe (return ()) f)
  
  newEmptyWBMVar :: (a -> IO ()) -> STM (WBMVar a)
  newEmptyWBMVar f = newWBVar Nothing (maybe (return ()) f)
  
  readWBMVar :: WBMVar a -> STM a
  readWBMVar v = readWBVar v >>= maybe retry return
  
  takeWBMVar :: WBMVar a -> IO a
  takeWBMVar v = atomically $ do
    a <- readWBMVar v
    writeTVar (wbValue v) Nothing
    return a
  
  writeWBMVar :: WBMVar a -> a -> Defer IO STM ()
  writeWBMVar v = writeWBVar v . Just
  
  withWBMVar :: WBMVar a -> (a -> IO b) -> IO b
  withWBMVar v =
    bracket (takeWBMVar v) (atomically . writeTVar (wbValue v) . Just)
    -- omit the write-back, we know the value didn't change.
  
  modifyWBMVar :: WBMVar a -> (a -> IO (a,b)) -> IO b
  modifyWBMVar v f =
    mask $ \restore -> do
      a <- takeWBMVar v
      (a',b) <- restore (f a)
        `onException` atomically (writeTVar (wbValue v) (Just a))
      immediately $ writeWBMVar v a'
      return b
  