@@ -14,13 +14,12 @@ module Control.Monad.Trans.Resource
1414import Control.Monad.Trans.Reader (ReaderT, runReaderT, ask)
1515import Data.IntMap (IntMap)
1616import qualified Data.IntMap as IntMap
17−import Data.Unique (newUnique, hashUnique)
1817import Control.Exception.Lifted (try, finally, SomeException, mask, mask_)
1918import Control.Monad.Trans.Control (MonadBaseControl)
2019import qualified Data.IORef as I
2120import Control.Monad.Base (MonadBase, liftBase)
2221
23−newtype ReleaseMap m = ReleaseMap (IntMap (ResourceT m ()))
22+data ReleaseMap m = ReleaseMap !Int !(IntMap (ResourceT m ()))
2423newtype ReleaseKey = ReleaseKey Int
2524
2625type ResourceT m = ReaderT (I.IORef (ReleaseMap m)) m
@@ -34,10 +33,10 @@ with acquire rel = mask_ $ do
3433 key <- register $ rel a
3534 return (key, a)
3635
37−modify :: MonadBase IO m => (ReleaseMap m -> ReleaseMap m) -> ResourceT m ()
36+modify :: MonadBase IO m => (ReleaseMap m -> (ReleaseMap m, a)) -> ResourceT m a
3837modify f = do
3938 istate <- ask
40− liftBase $ I.atomicModifyIORef istate $ \a -> (f a, ())
39+ liftBase $ I.atomicModifyIORef istate $ \a -> f a
4140
4241get :: MonadBase IO m => ResourceT m (ReleaseMap m)
4342get = ask >>= liftBase . I.readIORef
@@ -45,33 +44,32 @@ get = ask >>= liftBase . I.readIORef
4544register :: MonadBaseControl IO m
4645 => ResourceT m ()
4746 -> ResourceT m ReleaseKey
48−register rel = mask_ $ do
49− key <- liftBase $ fmap hashUnique newUnique
50− modify (insert key rel)
51− return (ReleaseKey key)
52− where
53− insert k v (ReleaseMap m) = ReleaseMap $ IntMap.insert k v m
47+register rel = mask_ $
48+ modify $ \(ReleaseMap key m) ->
49+ ( ReleaseMap (key + 1) (IntMap.insert key rel m)
50+ , ReleaseKey key
51+ )
5452
5553release :: MonadBaseControl IO m
5654 => ReleaseKey
5755 -> ResourceT m ()
5856release (ReleaseKey key) = mask $ \restore -> do
59− ReleaseMap m <- get
57+ ReleaseMap _ m <- get
6058 case IntMap.lookup key m of
6159 Nothing -> return () -- maybe we should throw an exception?
6260 Just r -> modify (delete key) >> restore (try' r >> return ())
6361 where
64− delete k (ReleaseMap m) = ReleaseMap $ IntMap.delete k m
62+ delete k (ReleaseMap next m) = (ReleaseMap next $ IntMap.delete k m, ())
6563
6664runResourceT :: MonadBaseControl IO m
6765 => ResourceT m a
6866 -> m a
6967runResourceT r = do
70− istate <- liftBase $ I.newIORef $ ReleaseMap IntMap.empty
68+ istate <- liftBase $ I.newIORef $ ReleaseMap minBound IntMap.empty
7169 runReaderT (r `finally` cleanup) istate
7270 where
7371 cleanup = do
74− ReleaseMap m <- get
72+ ReleaseMap _ m <- get
7573 if IntMap.null m
7674 then return ()
7775 else do