@@ -1,3 +1,4 @@
1+{-# LANGUAGE BangPatterns #-}
1-- | This module defines the 'WeakBag' type, which represents a mutable
2-- | This module defines the 'WeakBag' type, which represents a mutable
2-- collection of items that does not cause the items to be retained in memory.
3-- collection of items that does not cause the items to be retained in memory.
3-- This is useful for situations where a value needs to be inspected or modified
4-- This is useful for situations where a value needs to be inspected or modified
@@ -12,7 +13,6 @@ module Data.WeakBag
15−import Control.Concurrent.STM
16import Control.Exception
16import Control.Exception
17import Control.Monad hiding (forM_, mapM_)
17import Control.Monad hiding (forM_, mapM_)
18import Control.Monad.IO.Class
18import Control.Monad.IO.Class
@@ -28,8 +28,8 @@ import Prelude hiding (mapM_, traverse)
28-- that is, they can still be garbage-collected. As long as the @a@s remain
28-- that is, they can still be garbage-collected. As long as the @a@s remain
29-- alive, the 'WeakBag' will continue to refer to them.
29-- alive, the 'WeakBag' will continue to refer to them.
30data WeakBag a = WeakBag
30data WeakBag a = WeakBag
31− { _weakBag_nextId :: {-# UNPACK #-} !(TVar Int) --TODO: what if this wraps around?
31+ { _weakBag_nextId :: {-# UNPACK #-} !(IORef Int) --TODO: what if this wraps around?
32− , _weakBag_children :: {-# UNPACK #-} !(TVar (IntMap (Weak a)))
32+ , _weakBag_children :: {-# UNPACK #-} !(IORef (IntMap (Weak a)))
35-- | When inserting an item into a 'WeakBag', a 'WeakBagTicket' is returned. If
35-- | When inserting an item into a 'WeakBag', a 'WeakBagTicket' is returned. If
@@ -56,24 +56,19 @@ insert :: a -- ^ The item
56insert a (WeakBag nextId children) wbRef finalizer = {-# SCC "insert" #-} do
56insert a (WeakBag nextId children) wbRef finalizer = {-# SCC "insert" #-} do
57 a' <- evaluate a
57 a' <- evaluate a
58 wbRef' <- evaluate wbRef
58 wbRef' <- evaluate wbRef
59− myId <- atomically $ do
59+ myId <- atomicModifyIORef' nextId $ \n -> (succ n, n)
60− myId <- readTVar nextId
61− writeTVar nextId $! succ myId
63 let cleanup = do
60 let cleanup = do
64 wb <- readIORef wbRef'
61 wb <- readIORef wbRef'
65 mb <- deRefWeak wb
62 mb <- deRefWeak wb
66 forM_ mb $ \b -> do
63 forM_ mb $ \b -> do
67− isLastNode <- atomically $ do --TODO: Should this run even when mb is Nothing?
64+ csWithoutMe <- atomicModifyIORef children $ \cs ->
68− cs <- readTVar children
69− let csWithoutMe = IntMap.delete myId cs
65+ let !csWithoutMe = IntMap.delete myId cs
70− writeTVar children $! csWithoutMe
66+ in (csWithoutMe, csWithoutMe)
71− return $ IntMap.size csWithoutMe == 0
67+ when (IntMap.null csWithoutMe) $ finalizer b
72− when isLastNode $ finalizer b
75 wa <- mkWeakPtr a' $ Just cleanup
70 wa <- mkWeakPtr a' $ Just cleanup
76− atomically $ modifyTVar' children $ IntMap.insert myId wa
71+ atomicModifyIORef' children $ \cs -> (IntMap.insert myId wa cs, ())
77 return $ WeakBagTicket
72 return $ WeakBagTicket
78 { _weakBagTicket_weakItem = wa
73 { _weakBagTicket_weakItem = wa
79 , _weakBagTicket_item = a'
74 , _weakBagTicket_item = a'
@@ -83,8 +78,8 @@ insert a (WeakBag nextId children) wbRef finalizer = {-# SCC "insert" #-} do
83{-# INLINE empty #-}
78{-# INLINE empty #-}
84empty :: IO (WeakBag a)
79empty :: IO (WeakBag a)
85empty = {-# SCC "empty" #-} do
80empty = {-# SCC "empty" #-} do
86− nextId <- newTVarIO 1
81+ nextId <- newIORef 1
87− children <- newTVarIO IntMap.empty
82+ children <- newIORef IntMap.empty
88 let bag = WeakBag
83 let bag = WeakBag
89 { _weakBag_nextId = nextId
84 { _weakBag_nextId = nextId
90 , _weakBag_children = children
85 , _weakBag_children = children
@@ -107,7 +102,7 @@ singleton a wbRef finalizer = {-# SCC "singleton" #-} do
107-- is made about the order of the traversal.
102-- is made about the order of the traversal.
108traverse :: MonadIO m => WeakBag a -> (a -> m ()) -> m ()
103traverse :: MonadIO m => WeakBag a -> (a -> m ()) -> m ()
109traverse (WeakBag _ children) f = {-# SCC "traverse" #-} do
104traverse (WeakBag _ children) f = {-# SCC "traverse" #-} do
110− cs <- liftIO $ readTVarIO children
105+ cs <- liftIO $ readIORef children
111 forM_ cs $ \c -> do
106 forM_ cs $ \c -> do
112 ma <- liftIO $ deRefWeak c
107 ma <- liftIO $ deRefWeak c
113 mapM_ f ma
108 mapM_ f ma