git.iko.soy/reflex
on dan-perf help sign in
reflex/src/Data/WeakBag.hs
121 lines 4.6 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 ) where
21import Control.Exception
22import Control.Monad hiding (forM_, mapM_)
23import Control.Monad.IO.Class
24import Data.Foldable (forM_, mapM_)
25import Data.IntMap.Strict (IntMap)
26import qualified Data.IntMap.Strict as IntMap
27import Data.IORef
28import System.Mem.Weak
30import Prelude hiding (mapM_, traverse)
32-- | A @WeakBag a@ holds a set of values of type @a@, but does not retain them -
33-- that is, they can still be garbage-collected. As long as the @a@s remain
34-- alive, the 'WeakBag' will continue to refer to them.
35data WeakBag a = WeakBag
36 { _weakBag_nextId :: {-# UNPACK #-} !(IORef Int) --TODO: what if this wraps around?
37 , _weakBag_children :: {-# UNPACK #-} !(IORef (IntMap (Weak a)))
38 }
40-- | When inserting an item into a 'WeakBag', a 'WeakBagTicket' is returned. If
41-- the caller retains the ticket, the item is guranteed to stay in memory (and
42-- thus in the 'WeakBag'). The ticket can also be used to remove the item from
43-- the 'WeakBag' prematurely (i.e. while it is still alive), using 'remove'.
44data WeakBagTicket = forall a. WeakBagTicket
45 { _weakBagTicket_weakItem :: {-# UNPACK #-} !(Weak a)
46 , _weakBagTicket_item :: {-# NOUNPACK #-} !a
47 }
49-- | Insert an item into a 'WeakBag'.
50{-# INLINE insert #-}
51insert :: a -- ^ The item
52 -> WeakBag a -- ^ The 'WeakBag' to insert into
53 -> IORef (Weak b) -- ^ An arbitrary value to be used in the following
54 -- callback
55 -> (b -> IO ()) -- ^ A callback to be invoked when the item is removed
56 -- (whether automatically by the item being garbage
57 -- collected or manually via 'remove')
58 -> IO WeakBagTicket -- ^ Returns a 'WeakBagTicket' that ensures the item
59 -- is retained and allows the item to be removed.
60insert a (WeakBag nextId children) wbRef finalizer = {-# SCC "insert" #-} do
61 a' <- evaluate a
62 wbRef' <- evaluate wbRef
63 myId <- atomicModifyIORef' nextId $ \n -> (succ n, n)
64 let cleanup = do
65 wb <- readIORef wbRef'
66 mb <- deRefWeak wb
67 forM_ mb $ \b -> do
68 csWithoutMe <- atomicModifyIORef children $ \cs ->
69 let !csWithoutMe = IntMap.delete myId cs
70 in (csWithoutMe, csWithoutMe)
71 when (IntMap.null csWithoutMe) $ finalizer b
72 return ()
73 return ()
74 wa <- mkWeakPtr a' $ Just cleanup
75 atomicModifyIORef' children $ \cs -> (IntMap.insert myId wa cs, ())
76 return $ WeakBagTicket
77 { _weakBagTicket_weakItem = wa
78 , _weakBagTicket_item = a'
79 }
81-- | Create an empty 'WeakBag'.
82{-# INLINE empty #-}
83empty :: IO (WeakBag a)
84empty = {-# SCC "empty" #-} do
85 nextId <- newIORef 1
86 children <- newIORef IntMap.empty
87 let bag = WeakBag
88 { _weakBag_nextId = nextId
89 , _weakBag_children = children
90 }
91 return bag
93-- | Create a 'WeakBag' with one item; equivalent to creating the 'WeakBag' with
94-- 'empty', then using 'insert'.
95{-# INLINE singleton #-}
96singleton :: a -> IORef (Weak b) -> (b -> IO ()) -> IO (WeakBag a, WeakBagTicket)
97singleton a wbRef finalizer = {-# SCC "singleton" #-} do
98 bag <- empty
99 ticket <- insert a bag wbRef finalizer
100 return (bag, ticket)
102{-# INLINE traverse #-}
103-- | Visit every node in the given list. If new nodes are appended during the
104-- traversal, they will not be visited. Every live node that was in the list
105-- when the traversal began will be visited exactly once; however, no guarantee
106-- is made about the order of the traversal.
107traverse :: MonadIO m => WeakBag a -> (a -> m ()) -> m ()
108traverse (WeakBag _ children) f = {-# SCC "traverse" #-} do
109 cs <- liftIO $ readIORef children
110 forM_ cs $ \c -> do
111 ma <- liftIO $ deRefWeak c
112 mapM_ f ma
114-- | Remove an item from the 'WeakBag'; does nothing if invoked multiple times
115-- on the same 'WeakBagTicket'.
116{-# INLINE remove #-}
117remove :: WeakBagTicket -> IO ()
118remove (WeakBagTicket w _) = {-# SCC "remove" #-} finalize w
119--TODO: Should 'remove' also drop the reference to the item?
121--TODO: can/should we provide a null WeakBagTicket?