-- | 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
    (Cleanup
c, IO ()
clean) <- IO (Cleanup, IO ())
newCleanup
    Cleanup -> IO a
act Cleanup
c 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 <- forall a. a -> IO (IORef a)
newIORef forall a b. (a -> b) -> a -> b
$ Int -> HashMap Int (IO ()) -> S
S Int
0 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 = forall a. IO a -> IO a
uninterruptibleMask_ forall a b. (a -> b) -> a -> b
$ do
            HashMap Int (IO ())
items <- forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef S
ref forall a b. (a -> b) -> a -> b
$ \S
s -> (S
s{items :: HashMap Int (IO ())
items=forall k v. HashMap k v
Map.empty}, S -> HashMap Int (IO ())
items S
s)
            forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (forall a. Num a => a -> a
negate forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall a b. (a -> b) -> a -> b
$ forall k v. HashMap k v -> [(k, v)]
Map.toList HashMap Int (IO ())
items
    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 = forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef S
ref 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 forall a. Num a => a -> a -> a
+ Int
1) (forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
Map.insert Int
i IO ()
act 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) = forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef S
ref forall a b. (a -> b) -> a -> b
$ \S
s -> (S
s{items :: HashMap Int (IO ())
items = forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
Map.delete Int
i 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) = forall a. IO a -> IO a
uninterruptibleMask_ forall a b. (a -> b) -> a -> b
$ do
    Maybe (IO ())
undo <- forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef S
ref forall a b. (a -> b) -> a -> b
$ \S
s -> (S
s{items :: HashMap Int (IO ())
items = forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
Map.delete Int
i forall a b. (a -> b) -> a -> b
$ S -> HashMap Int (IO ())
items S
s}, forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup Int
i forall a b. (a -> b) -> a -> b
$ S -> HashMap Int (IO ())
items S
s)
    forall a. a -> Maybe a -> a
fromMaybe (forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) Maybe (IO ())
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 =
    forall a. IO a -> IO a
mask_ forall a b. (a -> b) -> a -> b
$ do
        a
v <- IO a
acquire
        Cleanup -> IO () -> IO ReleaseKey
register Cleanup
cleanup forall a b. (a -> b) -> a -> b
$ a -> IO ()
release a
v
        forall (f :: * -> *) a. Applicative f => a -> f a
pure a
v