-- | Code for ensuring cleanup actions are run.
module General.Cleanup(
    Cleanup, withCleanup, addCleanup, addCleanup_
    ) where

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


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

newtype Cleanup = Cleanup (IORef S)


-- | Run with some cleanup scope. Regardless of exceptions/threads, all 'addCleanup' actions
--   will be run by the time it exits. The 'addCleanup' actions will be run in reverse order.
withCleanup :: (Cleanup -> IO a) -> IO a
withCleanup :: (Cleanup -> IO a) -> IO a
withCleanup act :: Cleanup -> IO a
act = 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 0 HashMap Int (IO ())
forall k v. HashMap k v
Map.empty
    Cleanup -> IO a
act (IORef S -> Cleanup
Cleanup IORef S
ref) IO a -> IO () -> IO a
forall a b. IO a -> IO b -> IO a
`finally` Cleanup -> IO ()
runCleanup (IORef S -> Cleanup
Cleanup IORef S
ref)

-- | Run all the cleanup actions immediately. Done automatically by withCleanup
runCleanup :: Cleanup -> IO ()
runCleanup :: Cleanup -> IO ()
runCleanup (Cleanup ref :: IORef S
ref) = 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
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

-- | Add a cleanup action to a 'Cleanup' scope, returning a way to remove (but not perform) that action.
--   If not removed by the time 'withCleanup' terminates then the cleanup action will be run then.
addCleanup :: Cleanup -> IO () -> IO (IO ())
addCleanup :: Cleanup -> IO () -> IO (IO ())
addCleanup (Cleanup ref :: IORef S
ref) act :: IO ()
act = IORef S -> (S -> (S, IO ())) -> IO (IO ())
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef S
ref ((S -> (S, IO ())) -> IO (IO ()))
-> (S -> (S, IO ())) -> IO (IO ())
forall a b. (a -> b) -> a -> b
$ \s :: 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
+ 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)) (IO () -> (S, IO ())) -> IO () -> (S, IO ())
forall a b. (a -> b) -> a -> b
$
        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
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}, ())

addCleanup_ :: Cleanup -> IO () -> IO ()
-- we could avoid inserting into the Map, but we need to store the pairs anyway
-- to unregister them in order, so might as well keep it simple
addCleanup_ :: Cleanup -> IO () -> IO ()
addCleanup_ c :: Cleanup
c act :: IO ()
act = IO (IO ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (IO ()) -> IO ()) -> IO (IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ Cleanup -> IO () -> IO (IO ())
addCleanup Cleanup
c IO ()
act