1+ {-# LANGUAGE BangPatterns #-}
2+ {-# LANGUAGE ExistentialQuantification #-}
3+ {-# OPTIONS_GHC -fplugin=Reflex.Optimizer #-}
14-- | This module defines the 'WeakBag' type, which represents a mutable
25-- collection of items that does not cause the items to be retained in memory.
36-- This is useful for situations where a value needs to be inspected or modified
@@ -12,7 +15,6 @@ module Data.WeakBag
1215 , remove
1316 ) where
1417
15- import Control.Concurrent.STM
1618import Control.Exception
1719import Control.Monad hiding (forM_ , mapM_ )
1820import Control.Monad.IO.Class
@@ -28,15 +30,15 @@ import Prelude hiding (mapM_, traverse)
2830-- that is, they can still be garbage-collected. As long as the @a@s remain
2931-- alive, the 'WeakBag' will continue to refer to them.
3032data WeakBag a = WeakBag
31- { _weakBag_nextId :: {-# UNPACK #-} ! (TVar Int -- TODO: what if this wraps around?
32- , _weakBag_children :: {-# UNPACK #-} ! (TVar (IntMap (Weak a )))
33+ { _weakBag_nextId :: {-# UNPACK #-} ! (IORef Int -- TODO: what if this wraps around?
34+ , _weakBag_children :: {-# UNPACK #-} ! (IORef (IntMap (Weak a )))
3335 }
3436
3537-- | When inserting an item into a 'WeakBag', a 'WeakBagTicket' is returned. If
3638-- the caller retains the ticket, the item is guranteed to stay in memory (and
3739-- thus in the 'WeakBag'). The ticket can also be used to remove the item from
3840-- the 'WeakBag' prematurely (i.e. while it is still alive), using 'remove'.
39- data WeakBagTicket a = WeakBagTicket
41+ data WeakBagTicket = forall a . WeakBagTicket
4042 { _weakBagTicket_weakItem :: {-# UNPACK #-} ! (Weak a )
4143 , _weakBagTicket_item :: {-# NOUNPACK #-} ! a
4244 }
@@ -50,30 +52,24 @@ insert :: a -- ^ The item
5052 -> (b -> IO () ) -- ^ A callback to be invoked when the item is removed
5153 -- (whether automatically by the item being garbage
5254 -- collected or manually via 'remove')
53- -> IO WeakBagTicket a ) -- ^ Returns a 'WeakBagTicket' that ensures the
54- -- item is retained and allows the item to be
55- -- removed.
55+ -> IO WeakBagTicket -- ^ Returns a 'WeakBagTicket' that ensures the item
56+ -- is retained and allows the item to be removed.
5657insert a (WeakBag nextId children) wbRef finalizer = {-# SCC "insert" #-} do
5758 a' <- evaluate a
5859 wbRef' <- evaluate wbRef
59- myId <- atomically $ do
60- myId <- readTVar nextId
61- writeTVar nextId $! succ myId
62- return myId
60+ myId <- atomicModifyIORef' nextId $ \ n -> (succ n, n)
6361 let cleanup = do
6462 wb <- readIORef wbRef'
6563 mb <- deRefWeak wb
6664 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
65+ csWithoutMe <- atomicModifyIORef children $ \ cs ->
66+ let ! csWithoutMe = IntMap. delete myId cs
67+ in (csWithoutMe, csWithoutMe)
68+ when (IntMap. null csWithoutMe) $ finalizer b
7369 return ()
7470 return ()
7571 wa <- mkWeakPtr a' $ Just cleanup
76- atomically $ modifyTVar ' children $ IntMap. insert myId wa
72+ atomicModifyIORef ' children $ \ cs -> ( IntMap. insert myId wa cs, () )
7773 return $ WeakBagTicket
7874 { _weakBagTicket_weakItem = wa
7975 , _weakBagTicket_item = a'
@@ -83,8 +79,8 @@ insert a (WeakBag nextId children) wbRef finalizer = {-# SCC "insert" #-} do
8379{-# INLINE empty #-}
8480empty :: IO WeakBag a )
8581empty = {-# SCC "empty" #-} do
86- nextId <- newTVarIO 1
87- children <- newTVarIO IntMap. empty
82+ nextId <- newIORef 1
83+ children <- newIORef IntMap. empty
8884 let bag = WeakBag
8985 { _weakBag_nextId = nextId
9086 , _weakBag_children = children
@@ -94,7 +90,7 @@ empty = {-# SCC "empty" #-} do
9490-- | Create a 'WeakBag' with one item; equivalent to creating the 'WeakBag' with
9591-- 'empty', then using 'insert'.
9692{-# INLINE singleton #-}
97- singleton :: a -> IORef (Weak b ) -> (b -> IO () ) -> IO WeakBag a , WeakBagTicket a )
93+ singleton :: a -> IORef (Weak b ) -> (b -> IO () ) -> IO WeakBag a , WeakBagTicket )
9894singleton a wbRef finalizer = {-# SCC "singleton" #-} do
9995 bag <- empty
10096 ticket <- insert a bag wbRef finalizer
@@ -107,14 +103,16 @@ singleton a wbRef finalizer = {-# SCC "singleton" #-} do
107103-- is made about the order of the traversal.
108104traverse :: MonadIO m => WeakBag a -> (a -> m () ) -> m ()
109105traverse (WeakBag _ children) f = {-# SCC "traverse" #-} do
110- cs <- liftIO $ readTVarIO children
106+ cs <- liftIO $ readIORef children
111107 forM_ cs $ \ c -> do
112108 ma <- liftIO $ deRefWeak c
113109 mapM_ f ma
114110
115111-- | Remove an item from the 'WeakBag'; does nothing if invoked multiple times
116112-- on the same 'WeakBagTicket'.
117113{-# INLINE remove #-}
118- remove :: WeakBagTicket a -> IO ()
119- remove = {-# SCC "remove" #-} finalize . _weakBagTicket_weakItem
114+ remove :: WeakBagTicket -> IO ()
115+ remove ( WeakBagTicket w _) = {-# SCC "remove" #-} finalize w
120116-- TODO: Should 'remove' also drop the reference to the item?
117+
118+ -- TODO: can/should we provide a null WeakBagTicket?
0 commit comments