-- | Code for ensuring cleanup actions are run.
module General.Cleanup(
    Cleanup, newCleanup, withCleanup,
    register, release, allocate, unprotect
    ) where

import Control.Exception
import qualified Data.HashMap.Strict as Map
import Data.IORef
import Data.List.Extra
import Data.Maybe


data S = S
    {S -> Int
unique :: {-# UNPACK #-} !Int -- next index to be used to items
    ,S -> HashMap Int (IO ())
items :: !(Map.HashMap Int (IO ()))
    }

newtype Cleanup = Cleanup (IORef S)

data ReleaseKey = ReleaseKey (IORef S) {-# UNPACK #-} !Int


-- | Run with some cleanup scope. Regardless of exceptions/threads, all 'register' actions
--   will be run by the time it exits.
--   The 'register' actions will be run in reverse order, i.e. the last to be added will be run first.
withCleanup :: (Cleanup -> IO a) -> IO a
withCleanup :: forall a. (Cleanup -> IO a) -> IO a
withCleanup Cleanup -> IO a
act = do
    (c, clean) <- IO (Cleanup, IO ())
newCleanup
    act c `finally` clean

newCleanup :: IO (Cleanup, IO ())
newCleanup :: IO (Cleanup, IO ())
newCleanup = do
    ref <- S -> IO (IORef S)
forall a. a -> IO (IORef a)
newIORef (S -> IO (IORef S)) -> S -> IO (IORef S)
forall a b. (a -> b) -> a -> b
$ Int -> HashMap Int (IO ()) -> S
S Int
0 HashMap Int (IO ())
forall k v. HashMap k v
Map.empty
    -- important to use uninterruptibleMask_ otherwise in things like allocateThread
    -- we might end up being interrupted and failing to close down the thread
    -- e.g. see https://github.com/digital-asset/ghcide/issues/381
    -- note that packages like safe-exceptions also use uninterruptibleMask_
    let clean = IO () -> IO ()
forall a. IO a -> IO a
uninterruptibleMask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            items <- IORef S
-> (S -> (S, HashMap Int (IO ()))) -> IO (HashMap Int (IO ()))
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef S
ref ((S -> (S, HashMap Int (IO ()))) -> IO (HashMap Int (IO ())))
-> (S -> (S, HashMap Int (IO ()))) -> IO (HashMap Int (IO ()))
forall a b. (a -> b) -> a -> b
$ \S
s -> (S
s{items=Map.empty}, S -> HashMap Int (IO ())
items S
s)
            mapM_ snd $ sortOn (negate . fst) $ Map.toList items
    pure (Cleanup ref, clean)


register :: Cleanup -> IO () -> IO ReleaseKey
register :: Cleanup -> IO () -> IO ReleaseKey
register (Cleanup IORef S
ref) IO ()
act = IORef S -> (S -> (S, ReleaseKey)) -> IO ReleaseKey
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef S
ref ((S -> (S, ReleaseKey)) -> IO ReleaseKey)
-> (S -> (S, ReleaseKey)) -> IO ReleaseKey
forall a b. (a -> b) -> a -> b
$ \S
s -> let i :: Int
i = S -> Int
unique S
s in
    (Int -> HashMap Int (IO ()) -> S
S (S -> Int
unique S
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> IO () -> HashMap Int (IO ()) -> HashMap Int (IO ())
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
Map.insert Int
i IO ()
act (HashMap Int (IO ()) -> HashMap Int (IO ()))
-> HashMap Int (IO ()) -> HashMap Int (IO ())
forall a b. (a -> b) -> a -> b
$ S -> HashMap Int (IO ())
items S
s), IORef S -> Int -> ReleaseKey
ReleaseKey IORef S
ref Int
i)

unprotect :: ReleaseKey -> IO ()
unprotect :: ReleaseKey -> IO ()
unprotect (ReleaseKey IORef S
ref Int
i) = IORef S -> (S -> (S, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef S
ref ((S -> (S, ())) -> IO ()) -> (S -> (S, ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \S
s -> (S
s{items = Map.delete i $ items s}, ())

release :: ReleaseKey -> IO ()
release :: ReleaseKey -> IO ()
release (ReleaseKey IORef S
ref Int
i) = IO () -> IO ()
forall a. IO a -> IO a
uninterruptibleMask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    undo <- IORef S -> (S -> (S, Maybe (IO ()))) -> IO (Maybe (IO ()))
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef S
ref ((S -> (S, Maybe (IO ()))) -> IO (Maybe (IO ())))
-> (S -> (S, Maybe (IO ()))) -> IO (Maybe (IO ()))
forall a b. (a -> b) -> a -> b
$ \S
s -> (S
s{items = Map.delete i $ items s}, Int -> HashMap Int (IO ()) -> Maybe (IO ())
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup Int
i (HashMap Int (IO ()) -> Maybe (IO ()))
-> HashMap Int (IO ()) -> Maybe (IO ())
forall a b. (a -> b) -> a -> b
$ S -> HashMap Int (IO ())
items S
s)
    fromMaybe (pure ()) undo

allocate :: Cleanup -> IO a -> (a -> IO ()) -> IO a
allocate :: forall a. Cleanup -> IO a -> (a -> IO ()) -> IO a
allocate Cleanup
cleanup IO a
acquire a -> IO ()
release =
    IO a -> IO a
forall a. IO a -> IO a
mask_ (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ do
        v <- IO a
acquire
        register cleanup $ release v
        pure v