git.iko.soy/patch
on bd3f014 help sign in
patch/src/Data/WeakBag.hs
124 lines 4.8 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
3{-# LANGUAGE ExistentialQuantification #-}
4#ifdef USE_REFLEX_OPTIMIZER
5{-# OPTIONS_GHC -fplugin=Reflex.Optimizer #-}
6#endif
7-- | This module defines the 'WeakBag' type, which represents a mutable
8-- collection of items that does not cause the items to be retained in memory.
9-- This is useful for situations where a value needs to be inspected or modified
10-- if it is still alive, but can be ignored if it is dead.
11module Data.WeakBag
12 ( WeakBag
13 , WeakBagTicket
14 , empty
15 , singleton
16 , insert
17 , traverse
18 , remove
19 -- * Internal functions
20 -- These will not always be available.
21 , _weakBag_children --TODO: Don't export this
22 ) where
24import Control.Exception
25import Control.Monad hiding (forM_, mapM_)
26import Control.Monad.IO.Class
27import Data.Foldable (forM_, mapM_)
28import Data.IntMap.Strict (IntMap)
29import qualified Data.IntMap.Strict as IntMap
30import Data.IORef
31import System.Mem.Weak
33import Prelude hiding (mapM_, traverse)
35-- | A @WeakBag a@ holds a set of values of type @a@, but does not retain them -
36-- that is, they can still be garbage-collected. As long as the @a@s remain
37-- alive, the 'WeakBag' will continue to refer to them.
38data WeakBag a = WeakBag
39 { _weakBag_nextId :: {-# UNPACK #-} !(IORef Int) --TODO: what if this wraps around?
40 , _weakBag_children :: {-# UNPACK #-} !(IORef (IntMap (Weak a))) -- ^ The items referenced by the WeakBag
41 }
43-- | When inserting an item into a 'WeakBag', a 'WeakBagTicket' is returned. If
44-- the caller retains the ticket, the item is guranteed to stay in memory (and
45-- thus in the 'WeakBag'). The ticket can also be used to remove the item from
46-- the 'WeakBag' prematurely (i.e. while it is still alive), using 'remove'.
47data WeakBagTicket = forall a. WeakBagTicket
48 { _weakBagTicket_weakItem :: {-# UNPACK #-} !(Weak a)
49 , _weakBagTicket_item :: {-# NOUNPACK #-} !a
50 }
52-- | Insert an item into a 'WeakBag'.
53{-# INLINE insert #-}
54insert :: a -- ^ The item
55 -> WeakBag a -- ^ The 'WeakBag' to insert into
56 -> IORef (Weak b) -- ^ An arbitrary value to be used in the following
57 -- callback
58 -> (b -> IO ()) -- ^ A callback to be invoked when the item is removed
59 -- (whether automatically by the item being garbage
60 -- collected or manually via 'remove')
61 -> IO WeakBagTicket -- ^ Returns a 'WeakBagTicket' that ensures the item
62 -- is retained and allows the item to be removed.
63insert a (WeakBag nextId children) wbRef finalizer = {-# SCC "insert" #-} do
64 a' <- evaluate a
65 wbRef' <- evaluate wbRef
66 myId <- atomicModifyIORef' nextId $ \n -> (succ n, n)
67 let cleanup = do
68 wb <- readIORef wbRef'
69 mb <- deRefWeak wb
70 forM_ mb $ \b -> do
71 csWithoutMe <- atomicModifyIORef children $ \cs ->
72 let !csWithoutMe = IntMap.delete myId cs
73 in (csWithoutMe, csWithoutMe)
74 when (IntMap.null csWithoutMe) $ finalizer b
75 return ()
76 return ()
77 wa <- mkWeakPtr a' $ Just cleanup
78 atomicModifyIORef' children $ \cs -> (IntMap.insert myId wa cs, ())
79 return $ WeakBagTicket
80 { _weakBagTicket_weakItem = wa
81 , _weakBagTicket_item = a'
82 }
84-- | Create an empty 'WeakBag'.
85{-# INLINE empty #-}
86empty :: IO (WeakBag a)
87empty = {-# SCC "empty" #-} do
88 nextId <- newIORef 1
89 children <- newIORef IntMap.empty
90 let bag = WeakBag
91 { _weakBag_nextId = nextId
92 , _weakBag_children = children
93 }
94 return bag
96-- | Create a 'WeakBag' with one item; equivalent to creating the 'WeakBag' with
97-- 'empty', then using 'insert'.
98{-# INLINE singleton #-}
99singleton :: a -> IORef (Weak b) -> (b -> IO ()) -> IO (WeakBag a, WeakBagTicket)
100singleton a wbRef finalizer = {-# SCC "singleton" #-} do
101 bag <- empty
102 ticket <- insert a bag wbRef finalizer
103 return (bag, ticket)
105{-# INLINE traverse #-}
106-- | Visit every node in the given list. If new nodes are appended during the
107-- traversal, they will not be visited. Every live node that was in the list
108-- when the traversal began will be visited exactly once; however, no guarantee
109-- is made about the order of the traversal.
110traverse :: MonadIO m => WeakBag a -> (a -> m ()) -> m ()
111traverse (WeakBag _ children) f = {-# SCC "traverse" #-} do
112 cs <- liftIO $ readIORef children
113 forM_ cs $ \c -> do
114 ma <- liftIO $ deRefWeak c
115 mapM_ f ma
117-- | Remove an item from the 'WeakBag'; does nothing if invoked multiple times
118-- on the same 'WeakBagTicket'.
119{-# INLINE remove #-}
120remove :: WeakBagTicket -> IO ()
121remove (WeakBagTicket w _) = {-# SCC "remove" #-} finalize w
122--TODO: Should 'remove' also drop the reference to the item?
124--TODO: can/should we provide a null WeakBagTicket?