{-# LANGUAGE NamedFieldPuns #-}
-- | Module: Lifetimes.Rc
-- Description: Support for working with reference-counted resources.
--
-- Rather than associating a resource with one lifetime, a reference counted
-- resource associates each *reference* with a lifetime, and is released when
-- all references have expired.
module Lifetimes.Rc
    ( Rc
    , addRef
    , refCounted
    ) where

import Control.Concurrent.STM
import Lifetimes
import Zhp

-- | A resource which is managed by reference counting.
data Rc a = Rc
    { Rc a -> TVar Int
count   :: TVar Int
    , Rc a -> a
value   :: a
    , Rc a -> IO ()
cleanup :: IO ()
    }

-- | Acquire a new reference.
addRef :: Rc a -> Acquire a
addRef :: Rc a -> Acquire a
addRef Rc a
rc =
    IO a -> (a -> IO ()) -> Acquire a
forall a. IO a -> (a -> IO ()) -> Acquire a
mkAcquire
        (STM a -> IO a
forall a. STM a -> IO a
atomically (STM a -> IO a) -> STM a -> IO a
forall a b. (a -> b) -> a -> b
$ Rc a -> STM a
forall a. Rc a -> STM a
incRef Rc a
rc)
        (\a
_ -> IO (IO ()) -> IO ()
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (IO (IO ()) -> IO ()) -> IO (IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ STM (IO ()) -> IO (IO ())
forall a. STM a -> IO a
atomically (STM (IO ()) -> IO (IO ())) -> STM (IO ()) -> IO (IO ())
forall a b. (a -> b) -> a -> b
$ Rc a -> STM (IO ())
forall a. Rc a -> STM (IO ())
decRef Rc a
rc)

resourceToRc :: Resource a -> STM (Rc a)
resourceToRc :: Resource a -> STM (Rc a)
resourceToRc Resource a
res = do
    a
value <- Resource a -> STM a
forall (m :: * -> *) a. MonadSTM m => Resource a -> m a
mustGetResource Resource a
res
    IO ()
cleanup <- Resource a -> STM (IO ())
forall (m :: * -> *) a. MonadSTM m => Resource a -> m (IO ())
detach Resource a
res
    TVar Int
count <- Int -> STM (TVar Int)
forall a. a -> STM (TVar a)
newTVar Int
1
    Rc a -> STM (Rc a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Rc :: forall a. TVar Int -> a -> IO () -> Rc a
Rc { TVar Int
count :: TVar Int
count :: TVar Int
count, IO ()
cleanup :: IO ()
cleanup :: IO ()
cleanup, a
value :: a
value :: a
value }


-- | Acquire a resource using refcounting. Takes an 'Acquire' for the underlying
-- resource, and returns one that acquires an initial reference to it. Additional
-- references may be created using 'addRef', and the underlying resource will be
-- kept alive until all resources are released.
refCounted :: Acquire a -> Acquire (Rc a)
refCounted :: Acquire a -> Acquire (Rc a)
refCounted Acquire a
acq = do
    Lifetime
lt <- Acquire Lifetime
currentLifetime
    IO (Rc a) -> Acquire (Rc a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Rc a) -> Acquire (Rc a)) -> IO (Rc a) -> Acquire (Rc a)
forall a b. (a -> b) -> a -> b
$ (Lifetime -> IO (Rc a)) -> IO (Rc a)
forall a. (Lifetime -> IO a) -> IO a
withLifetime ((Lifetime -> IO (Rc a)) -> IO (Rc a))
-> (Lifetime -> IO (Rc a)) -> IO (Rc a)
forall a b. (a -> b) -> a -> b
$ \Lifetime
tmpLt -> do
        Resource a
res <- Lifetime -> Acquire a -> IO (Resource a)
forall a. Lifetime -> Acquire a -> IO (Resource a)
acquire Lifetime
tmpLt Acquire a
acq
        Lifetime -> Acquire (Rc a) -> IO (Rc a)
forall a. Lifetime -> Acquire a -> IO a
acquireValue Lifetime
lt (Acquire (Rc a) -> IO (Rc a)) -> Acquire (Rc a) -> IO (Rc a)
forall a b. (a -> b) -> a -> b
$ IO (Rc a) -> (Rc a -> IO ()) -> Acquire (Rc a)
forall a. IO a -> (a -> IO ()) -> Acquire a
mkAcquire
            (STM (Rc a) -> IO (Rc a)
forall a. STM a -> IO a
atomically (STM (Rc a) -> IO (Rc a)) -> STM (Rc a) -> IO (Rc a)
forall a b. (a -> b) -> a -> b
$ Resource a -> STM (Rc a)
forall a. Resource a -> STM (Rc a)
resourceToRc Resource a
res)
            (IO (IO ()) -> IO ()
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (IO (IO ()) -> IO ()) -> (Rc a -> IO (IO ())) -> Rc a -> IO ()
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. STM (IO ()) -> IO (IO ())
forall a. STM a -> IO a
atomically (STM (IO ()) -> IO (IO ()))
-> (Rc a -> STM (IO ())) -> Rc a -> IO (IO ())
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Rc a -> STM (IO ())
forall a. Rc a -> STM (IO ())
decRef)


incRef :: Rc a -> STM a
incRef :: Rc a -> STM a
incRef Rc{TVar Int
count :: TVar Int
count :: forall a. Rc a -> TVar Int
count, a
value :: a
value :: forall a. Rc a -> a
value} = do
    TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
count Int -> Int
forall a. Enum a => a -> a
succ
    a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
value

decRef :: Rc a -> STM (IO ())
decRef :: Rc a -> STM (IO ())
decRef Rc{TVar Int
count :: TVar Int
count :: forall a. Rc a -> TVar Int
count, IO ()
cleanup :: IO ()
cleanup :: forall a. Rc a -> IO ()
cleanup} = do
    TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
count Int -> Int
forall a. Enum a => a -> a
pred
    Int
c <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
count
    IO () -> STM (IO ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IO () -> STM (IO ())) -> IO () -> STM (IO ())
forall a b. (a -> b) -> a -> b
$ case Int
c of
        Int
0 -> IO ()
cleanup
        Int
_ -> () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()