-- | 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 :: (Cleanup -> IO a) -> IO a
withCleanup Cleanup -> IO a
act = do
    (Cleanup
c, IO ()
clean) <- IO (Cleanup, IO ())
newCleanup
    Cleanup -> IO a
act Cleanup
c IO a -> IO () -> IO a
forall a b. IO a -> IO b -> IO a
`finally` IO ()
clean

newCleanup :: IO (Cleanup, IO ())
newCleanup :: IO (Cleanup, IO ())
newCleanup = do
    IORef S
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 ()
clean = IO () -> IO ()
forall a. IO a -> IO a
uninterruptibleMask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            HashMap Int (IO ())
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 :: HashMap Int (IO ())
items=HashMap Int (IO ())
forall k v. HashMap k v
Map.empty}, S -> HashMap Int (IO ())
items S
s)
            ((Int, IO ()) -> IO ()) -> [(Int, IO ())] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Int, IO ()) -> IO ()
forall a b. (a, b) -> b
snd ([(Int, IO ())] -> IO ()) -> [(Int, IO ())] -> IO ()
forall a b. (a -> b) -> a -> b
$ ((Int, IO ()) -> Int) -> [(Int, IO ())] -> [(Int, IO ())]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Int -> Int
forall a. Num a => a -> a
negate (Int -> Int) -> ((Int, IO ()) -> Int) -> (Int, IO ()) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, IO ()) -> Int
forall a b. (a, b) -> a
fst) ([(Int, IO ())] -> [(Int, IO ())])
-> [(Int, IO ())] -> [(Int, IO ())]
forall a b. (a -> b) -> a -> b
$ HashMap Int (IO ()) -> [(Int, IO ())]
forall k v. HashMap k v -> [(k, v)]
Map.toList HashMap Int (IO ())
items
    (Cleanup, IO ()) -> IO (Cleanup, IO ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IORef S -> Cleanup
Cleanup IORef S
ref, IO ()
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 :: HashMap Int (IO ())
items = Int -> HashMap Int (IO ()) -> HashMap Int (IO ())
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
Map.delete Int
i (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}, ())

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
    Maybe (IO ())
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 :: HashMap Int (IO ())
items = Int -> HashMap Int (IO ()) -> HashMap Int (IO ())
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
Map.delete Int
i (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}, 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)
    IO () -> Maybe (IO ()) -> IO ()
forall a. a -> Maybe a -> a
fromMaybe (() -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) Maybe (IO ())
undo

allocate :: Cleanup -> IO a -> (a -> IO ()) -> IO a
allocate :: 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
        a
v <- IO a
acquire
        Cleanup -> IO () -> IO ReleaseKey
register Cleanup
cleanup (IO () -> IO ReleaseKey) -> IO () -> IO ReleaseKey
forall a b. (a -> b) -> a -> b
$ a -> IO ()
release a
v
        a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
v