@@ -1,3 +1,4 @@
1+{-# LANGUAGE BangPatterns #-}
12-- | This module defines the 'WeakBag' type, which represents a mutable
23-- collection of items that does not cause the items to be retained in memory.
34-- This is useful for situations where a value needs to be inspected or modified
@@ -12,7 +13,6 @@ module Data.WeakBag
1213 , remove
1314 ) where
1415
15−import Control.Concurrent.STM
1616import Control.Exception
1717import Control.Monad hiding (forM_, mapM_)
1818import Control.Monad.IO.Class
@@ -28,8 +28,8 @@ import Prelude hiding (mapM_, traverse)
2828-- that is, they can still be garbage-collected. As long as the @a@s remain
2929-- alive, the 'WeakBag' will continue to refer to them.
3030data WeakBag a = WeakBag
31− { _weakBag_nextId :: {-# UNPACK #-} !(TVar Int) --TODO: what if this wraps around?
32− , _weakBag_children :: {-# UNPACK #-} !(TVar (IntMap (Weak a)))
31+ { _weakBag_nextId :: {-# UNPACK #-} !(IORef Int) --TODO: what if this wraps around?
32+ , _weakBag_children :: {-# UNPACK #-} !(IORef (IntMap (Weak a)))
3333 }
3434
3535-- | When inserting an item into a 'WeakBag', a 'WeakBagTicket' is returned. If
@@ -56,24 +56,19 @@ insert :: a -- ^ The item
5656insert a (WeakBag nextId children) wbRef finalizer = {-# SCC "insert" #-} do
5757 a' <- evaluate a
5858 wbRef' <- evaluate wbRef
59− myId <- atomically $ do
60− myId <- readTVar nextId
61− writeTVar nextId $! succ myId
62− return myId
59+ myId <- atomicModifyIORef' nextId $ \n -> (succ n, n)
6360 let cleanup = do
6461 wb <- readIORef wbRef'
6562 mb <- deRefWeak wb
6663 forM_ mb $ \b -> do
67− isLastNode <- atomically $ do --TODO: Should this run even when mb is Nothing?
68− cs <- readTVar children
69− let csWithoutMe = IntMap.delete myId cs
70− writeTVar children $! csWithoutMe
71− return $ IntMap.size csWithoutMe == 0
72− when isLastNode $ finalizer b
64+ csWithoutMe <- atomicModifyIORef children $ \cs ->
65+ let !csWithoutMe = IntMap.delete myId cs
66+ in (csWithoutMe, csWithoutMe)
67+ when (IntMap.null csWithoutMe) $ finalizer b
7368 return ()
7469 return ()
7570 wa <- mkWeakPtr a' $ Just cleanup
76− atomically $ modifyTVar' children $ IntMap.insert myId wa
71+ atomicModifyIORef' children $ \cs -> (IntMap.insert myId wa cs, ())
7772 return $ WeakBagTicket
7873 { _weakBagTicket_weakItem = wa
7974 , _weakBagTicket_item = a'
@@ -83,8 +78,8 @@ insert a (WeakBag nextId children) wbRef finalizer = {-# SCC "insert" #-} do
8378{-# INLINE empty #-}
8479empty :: IO (WeakBag a)
8580empty = {-# SCC "empty" #-} do
86− nextId <- newTVarIO 1
87− children <- newTVarIO IntMap.empty
81+ nextId <- newIORef 1
82+ children <- newIORef IntMap.empty
8883 let bag = WeakBag
8984 { _weakBag_nextId = nextId
9085 , _weakBag_children = children
@@ -107,7 +102,7 @@ singleton a wbRef finalizer = {-# SCC "singleton" #-} do
107102-- is made about the order of the traversal.
108103traverse :: MonadIO m => WeakBag a -> (a -> m ()) -> m ()
109104traverse (WeakBag _ children) f = {-# SCC "traverse" #-} do
110− cs <- liftIO $ readTVarIO children
105+ cs <- liftIO $ readIORef children
111106 forM_ cs $ \c -> do
112107 ma <- liftIO $ deRefWeak c
113108 mapM_ f ma