git.iko.soy/patch
on 52ca995 help sign in
patch/src/Data/WeakBag.hs
120 lines 4.7 kB history raw
1-- | 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-- This is useful for situations where a value needs to be inspected or modified
4-- if it is still alive, but can be ignored if it is dead.
5module Data.WeakBag
6 ( WeakBag
7 , WeakBagTicket
8 , empty
9 , singleton
10 , insert
11 , traverse
12 , remove
13 ) where
15import Control.Concurrent.STM
16import Control.Exception
17import Control.Monad hiding (forM_, mapM_)
18import Control.Monad.IO.Class
19import Data.Foldable (forM_, mapM_)
20import Data.IntMap.Strict (IntMap)
21import qualified Data.IntMap.Strict as IntMap
22import Data.IORef
23import System.Mem.Weak
25import Prelude hiding (mapM_, traverse)
27-- | A @WeakBag a@ holds a set of values of type @a@, but does not retain them -
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.
30data WeakBag a = WeakBag
31 { _weakBag_nextId :: {-# UNPACK #-} !(TVar Int) --TODO: what if this wraps around?
32 , _weakBag_children :: {-# UNPACK #-} !(TVar (IntMap (Weak a)))
33 }
35-- | When inserting an item into a 'WeakBag', a 'WeakBagTicket' is returned. If
36-- the caller retains the ticket, the item is guranteed to stay in memory (and
37-- thus in the 'WeakBag'). The ticket can also be used to remove the item from
38-- the 'WeakBag' prematurely (i.e. while it is still alive), using 'remove'.
39data WeakBagTicket a = WeakBagTicket
40 { _weakBagTicket_weakItem :: {-# UNPACK #-} !(Weak a)
41 , _weakBagTicket_item :: {-# NOUNPACK #-} !a
42 }
44-- | Insert an item into a 'WeakBag'.
45{-# INLINE insert #-}
46insert :: a -- ^ The item
47 -> WeakBag a -- ^ The 'WeakBag' to insert into
48 -> IORef (Weak b) -- ^ An arbitrary value to be used in the following
49 -- callback
50 -> (b -> IO ()) -- ^ A callback to be invoked when the item is removed
51 -- (whether automatically by the item being garbage
52 -- 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.
56insert a (WeakBag nextId children) wbRef finalizer = {-# SCC "insert" #-} do
57 a' <- evaluate a
58 wbRef' <- evaluate wbRef
59 myId <- atomically $ do
60 myId <- readTVar nextId
61 writeTVar nextId $! succ myId
62 return myId
63 let cleanup = do
64 wb <- readIORef wbRef'
65 mb <- deRefWeak wb
66 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
73 return ()
74 return ()
75 wa <- mkWeakPtr a' $ Just cleanup
76 atomically $ modifyTVar' children $ IntMap.insert myId wa
77 return $ WeakBagTicket
78 { _weakBagTicket_weakItem = wa
79 , _weakBagTicket_item = a'
80 }
82-- | Create an empty 'WeakBag'.
83{-# INLINE empty #-}
84empty :: IO (WeakBag a)
85empty = {-# SCC "empty" #-} do
86 nextId <- newTVarIO 1
87 children <- newTVarIO IntMap.empty
88 let bag = WeakBag
89 { _weakBag_nextId = nextId
90 , _weakBag_children = children
91 }
92 return bag
94-- | Create a 'WeakBag' with one item; equivalent to creating the 'WeakBag' with
95-- 'empty', then using 'insert'.
96{-# INLINE singleton #-}
97singleton :: a -> IORef (Weak b) -> (b -> IO ()) -> IO (WeakBag a, WeakBagTicket a)
98singleton a wbRef finalizer = {-# SCC "singleton" #-} do
99 bag <- empty
100 ticket <- insert a bag wbRef finalizer
101 return (bag, ticket)
103{-# INLINE traverse #-}
104-- | Visit every node in the given list. If new nodes are appended during the
105-- traversal, they will not be visited. Every live node that was in the list
106-- when the traversal began will be visited exactly once; however, no guarantee
107-- is made about the order of the traversal.
108traverse :: MonadIO m => WeakBag a -> (a -> m ()) -> m ()
109traverse (WeakBag _ children) f = {-# SCC "traverse" #-} do
110 cs <- liftIO $ readTVarIO children
111 forM_ cs $ \c -> do
112 ma <- liftIO $ deRefWeak c
113 mapM_ f ma
115-- | Remove an item from the 'WeakBag'; does nothing if invoked multiple times
116-- on the same 'WeakBagTicket'.
117{-# INLINE remove #-}
118remove :: WeakBagTicket a -> IO ()
119remove = {-# SCC "remove" #-} finalize . _weakBagTicket_weakItem
120--TODO: Should 'remove' also drop the reference to the item?