{-| Lightweight reference counter that runs a release action when the count
hits zero. Useful for keeping a Vulkan object alive across an unknown number
of in-flight frames — bump the count when a frame starts using it, drop the
count when the frame retires, and the object is destroyed promptly after the
last frame finishes.
-}
module Vulkan.Utils.RefCounted
  ( RefCounted
  , newRefCounted
  , releaseRefCounted
  , takeRefCounted
  , resourceTRefCount
  ) where

import Control.Exception (mask, throwIO)
import Control.Monad
import Control.Monad.IO.Class (MonadIO (..))
import Control.Monad.Trans.Resource (MonadResource, allocate_)
import Data.IORef
import GHC.IO.Exception (IOErrorType (UserError), IOException (IOError))

-- | A 'RefCounted' will perform the specified action when the count reaches 0
data RefCounted = RefCounted
  { RefCounted -> IORef Int
rcCount :: IORef Int
  , RefCounted -> IO ()
rcAction :: IO ()
  }

-- | Create a counter with a value of 1
newRefCounted :: (MonadIO m) => IO () -> m RefCounted
newRefCounted :: forall (m :: * -> *). MonadIO m => IO () -> m RefCounted
newRefCounted IO ()
rcAction = do
  rcCount <- IO (IORef Int) -> m (IORef Int)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IORef Int) -> m (IORef Int))
-> IO (IORef Int) -> m (IORef Int)
forall a b. (a -> b) -> a -> b
$ Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
1
  pure RefCounted{..}

{- | Decrement the value, the action will be run promptly and in
this thread if the counter reached 0.
-}
releaseRefCounted :: (MonadIO m) => RefCounted -> m ()
releaseRefCounted :: forall (m :: * -> *). MonadIO m => RefCounted -> m ()
releaseRefCounted RefCounted{IO ()
IORef Int
rcCount :: RefCounted -> IORef Int
rcAction :: RefCounted -> IO ()
rcCount :: IORef Int
rcAction :: IO ()
..} = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO ()) -> IO ())
-> ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
_ ->
  IORef Int -> (Int -> (Int, Int)) -> IO Int
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Int
rcCount (\Int
c -> (Int -> Int
forall a. Enum a => a -> a
pred Int
c, Int -> Int
forall a. Enum a => a -> a
pred Int
c)) IO Int -> (Int -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Int
0 -> IO ()
rcAction
    Int
n
      | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 ->
          IO () -> IO ()
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> IO ()) -> (IOException -> IO ()) -> IOException -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IOException -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (IOException -> IO ()) -> IOException -> IO ()
forall a b. (a -> b) -> a -> b
$
            Maybe Handle
-> IOErrorType
-> String
-> String
-> Maybe CInt
-> Maybe String
-> IOException
IOError
              Maybe Handle
forall a. Maybe a
Nothing
              IOErrorType
UserError
              String
""
              String
"Ref counted value decremented below 0"
              Maybe CInt
forall a. Maybe a
Nothing
              Maybe String
forall a. Maybe a
Nothing
    Int
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Increment the counter by 1
takeRefCounted :: (MonadIO m) => RefCounted -> m ()
takeRefCounted :: forall (m :: * -> *). MonadIO m => RefCounted -> m ()
takeRefCounted RefCounted{IO ()
IORef Int
rcCount :: RefCounted -> IORef Int
rcAction :: RefCounted -> IO ()
rcCount :: IORef Int
rcAction :: IO ()
..} =
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef Int -> (Int -> (Int, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Int
rcCount (\Int
c -> (Int -> Int
forall a. Enum a => a -> a
succ Int
c, ()))

-- | Hold a reference for the duration of the 'MonadResource' action
resourceTRefCount :: (MonadResource f) => RefCounted -> f ()
resourceTRefCount :: forall (f :: * -> *). MonadResource f => RefCounted -> f ()
resourceTRefCount RefCounted
r = f ReleaseKey -> f ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (f ReleaseKey -> f ()) -> f ReleaseKey -> f ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO () -> f ReleaseKey
forall (m :: * -> *) a.
MonadResource m =>
IO a -> IO () -> m ReleaseKey
allocate_ (RefCounted -> IO ()
forall (m :: * -> *). MonadIO m => RefCounted -> m ()
takeRefCounted RefCounted
r) (RefCounted -> IO ()
forall (m :: * -> *). MonadIO m => RefCounted -> m ()
releaseRefCounted RefCounted
r)