module General.Cleanup(
Cleanup, withCleanup, addCleanup
) where
import Control.Exception as E
import Control.Monad
import qualified Data.HashMap.Strict as Map
import Data.Function
import Data.IORef
import Data.List
data S = S {unique :: !Int, items :: Map.HashMap Int (IO ())}
newtype Cleanup = Cleanup (IORef S)
withCleanup :: (Cleanup -> IO a) -> IO a
withCleanup act = do
ref <- newIORef $ S 0 Map.empty
act (Cleanup ref) `finally` do
items <- atomicModifyIORef ref $ \s -> (s{items=Map.empty}, items s)
mapM_ snd $ sortBy (compare `on` negate . fst) $ Map.toList items
addCleanup :: Cleanup -> IO () -> IO (Bool -> IO ())
addCleanup (Cleanup ref) act = atomicModifyIORef ref $ \s -> let i = unique s in
(,) (S (unique s + 1) (Map.insert i act $ items s)) $ \b ->
join $ atomicModifyIORef ref $ \s -> case Map.lookup i $ items s of
Nothing -> (s, return ())
Just act -> (s{items = Map.delete i $ items s}, when b act)