-- | 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 {unique :: {-# UNPACK #-} !Int, 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 act = do
    ref <- newIORef $ S 0 Map.empty
    act (Cleanup ref) `finally` runCleanup (Cleanup ref)

-- | Run all the cleanup actions immediately. Done automatically by withCleanup
runCleanup :: Cleanup -> IO ()
runCleanup (Cleanup ref) = do
    items <- atomicModifyIORef' ref $ \s -> (s{items=Map.empty}, items s)
    mapM_ snd $ sortOn (negate . fst) $ Map.toList 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 ref) act = atomicModifyIORef' ref $ \s -> let i = unique s in
    (,) (S (unique s + 1) (Map.insert i act $ items s)) $
        atomicModifyIORef' ref $ \s -> (s{items = Map.delete i $ items 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_ c act = void $ addCleanup c act