{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE ExistentialQuantification #-}
#ifdef USE_REFLEX_OPTIMIZER
{-# OPTIONS_GHC -fplugin=Reflex.Optimizer #-}
#endif
-- | This module defines the 'WeakBag' type, which represents a mutable
-- collection of items that does not cause the items to be retained in memory.
-- This is useful for situations where a value needs to be inspected or modified
-- if it is still alive, but can be ignored if it is dead.
module Data.WeakBag
  ( WeakBag
  , WeakBagTicket
  , empty
  , singleton
  , insert
  , traverse
  , traverse_
  , remove
  -- * Internal functions
  -- These will not always be available.
  , _weakBag_children --TODO: Don't export this
  ) where

import Prelude hiding (traverse)

import Control.Exception
import Control.Monad
import Control.Monad.IO.Class
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import Data.IORef
import System.Mem.Weak

-- | A 'WeakBag' holds a set of values of type @/a/@, but does not retain them -
-- that is, they can still be garbage-collected.  As long as the @/a/@ values remain
-- alive, the 'WeakBag' will continue to refer to them.
data WeakBag a = WeakBag
  { forall a. WeakBag a -> IORef Int
_weakBag_nextId :: {-# UNPACK #-} !(IORef Int) --TODO: what if this wraps around?
  , forall a. WeakBag a -> IORef (IntMap (Weak a))
_weakBag_children :: {-# UNPACK #-} !(IORef (IntMap (Weak a))) -- ^ The items referenced by the WeakBag
  }

-- | When inserting an item into a 'WeakBag', a 'WeakBagTicket' is returned.  If
-- the caller retains the ticket, the item is guranteed to stay in memory (and
-- thus in the 'WeakBag').  The ticket can also be used to remove the item from
-- the 'WeakBag' prematurely (i.e. while it is still alive), using 'remove'.
data WeakBagTicket = forall a. WeakBagTicket
  { ()
_weakBagTicket_weakItem :: {-# UNPACK #-} !(Weak a)
  , ()
_weakBagTicket_item :: {-# NOUNPACK #-} !a
  }

-- | Insert an item into a 'WeakBag'.
{-# INLINE insert #-}
insert :: a -- ^ The item
       -> WeakBag a -- ^ The 'WeakBag' to insert into
       -> IORef (Weak b) -- ^ An arbitrary value to be used in the following
                         -- callback
       -> (b -> IO ()) -- ^ A callback to be invoked when the item is removed
                       -- (whether automatically by the item being garbage
                       -- collected or manually via 'remove')
       -> IO WeakBagTicket -- ^ Returns a 'WeakBagTicket' that ensures the item
                           -- is retained and allows the item to be removed.
insert :: forall a b.
a
-> WeakBag a -> IORef (Weak b) -> (b -> IO ()) -> IO WeakBagTicket
insert a
a (WeakBag IORef Int
nextId IORef (IntMap (Weak a))
children) IORef (Weak b)
wbRef b -> IO ()
finalizer = {-# SCC "insert" #-} do
  a' <- a -> IO a
forall a. a -> IO a
evaluate a
a
  wbRef' <- evaluate wbRef
  myId <- atomicModifyIORef' nextId $ \Int
n -> (Int -> Int
forall a. Enum a => a -> a
succ Int
n, Int
n)
  let cleanup = do
        wb <- IORef (Weak b) -> IO (Weak b)
forall a. IORef a -> IO a
readIORef IORef (Weak b)
wbRef'
        mb <- deRefWeak wb
        forM_ mb $ \b
b -> do
          csWithoutMe <- IORef (IntMap (Weak a))
-> (IntMap (Weak a) -> (IntMap (Weak a), IntMap (Weak a)))
-> IO (IntMap (Weak a))
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef (IntMap (Weak a))
children ((IntMap (Weak a) -> (IntMap (Weak a), IntMap (Weak a)))
 -> IO (IntMap (Weak a)))
-> (IntMap (Weak a) -> (IntMap (Weak a), IntMap (Weak a)))
-> IO (IntMap (Weak a))
forall a b. (a -> b) -> a -> b
$ \IntMap (Weak a)
cs ->
            let !csWithoutMe :: IntMap (Weak a)
csWithoutMe = Int -> IntMap (Weak a) -> IntMap (Weak a)
forall a. Int -> IntMap a -> IntMap a
IntMap.delete Int
myId IntMap (Weak a)
cs
            in (IntMap (Weak a)
csWithoutMe, IntMap (Weak a)
csWithoutMe)
          when (IntMap.null csWithoutMe) $ finalizer b
  wa <- mkWeakPtr a' $ Just cleanup
  atomicModifyIORef' children $ \IntMap (Weak a)
cs -> (Int -> Weak a -> IntMap (Weak a) -> IntMap (Weak a)
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
myId Weak a
wa IntMap (Weak a)
cs, ())
  return $ WeakBagTicket
    { _weakBagTicket_weakItem = wa
    , _weakBagTicket_item = a'
    }

-- | Create an empty 'WeakBag'.
{-# INLINE empty #-}
empty :: IO (WeakBag a)
empty :: forall a. IO (WeakBag a)
empty = {-# SCC "empty" #-} do
  nextId <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
1
  children <- newIORef IntMap.empty
  let bag = WeakBag
        { _weakBag_nextId :: IORef Int
_weakBag_nextId = IORef Int
nextId
        , _weakBag_children :: IORef (IntMap (Weak a))
_weakBag_children = IORef (IntMap (Weak a))
children
        }
  return bag

-- | Create a 'WeakBag' with one item; equivalent to creating the 'WeakBag' with
-- 'empty', then using 'insert'.
{-# INLINE singleton #-}
singleton :: a -> IORef (Weak b) -> (b -> IO ()) -> IO (WeakBag a, WeakBagTicket)
singleton :: forall a b.
a
-> IORef (Weak b) -> (b -> IO ()) -> IO (WeakBag a, WeakBagTicket)
singleton a
a IORef (Weak b)
wbRef b -> IO ()
finalizer = {-# SCC "singleton" #-} do
  bag <- IO (WeakBag a)
forall a. IO (WeakBag a)
empty
  ticket <- insert a bag wbRef finalizer
  return (bag, ticket)

{-# INLINE traverse_ #-}
-- | Visit every node in the given list.  If new nodes are appended during the
-- traversal, they will not be visited.  Every live node that was in the list
-- when the traversal began will be visited exactly once; however, no guarantee
-- is made about the order of the traversal.
traverse_ :: MonadIO m => WeakBag a -> (a -> m ()) -> m ()
traverse_ :: forall (m :: * -> *) a.
MonadIO m =>
WeakBag a -> (a -> m ()) -> m ()
traverse_ (WeakBag IORef Int
_ IORef (IntMap (Weak a))
children) a -> m ()
f = {-# SCC "traverse" #-} do
  cs <- IO (IntMap (Weak a)) -> m (IntMap (Weak a))
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IntMap (Weak a)) -> m (IntMap (Weak a)))
-> IO (IntMap (Weak a)) -> m (IntMap (Weak a))
forall a b. (a -> b) -> a -> b
$ IORef (IntMap (Weak a)) -> IO (IntMap (Weak a))
forall a. IORef a -> IO a
readIORef IORef (IntMap (Weak a))
children
  forM_ cs $ \Weak a
c -> do
    ma <- IO (Maybe a) -> m (Maybe a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe a) -> m (Maybe a)) -> IO (Maybe a) -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ Weak a -> IO (Maybe a)
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak a
c
    mapM_ f ma

{-# DEPRECATED traverse "Use 'traverse_' instead" #-}
traverse :: MonadIO m => WeakBag a -> (a -> m ()) -> m ()
traverse :: forall (m :: * -> *) a.
MonadIO m =>
WeakBag a -> (a -> m ()) -> m ()
traverse = WeakBag a -> (a -> m ()) -> m ()
forall (m :: * -> *) a.
MonadIO m =>
WeakBag a -> (a -> m ()) -> m ()
traverse_

-- | Remove an item from the 'WeakBag'; does nothing if invoked multiple times
-- on the same 'WeakBagTicket'.
{-# INLINE remove #-}
remove :: WeakBagTicket -> IO ()
remove :: WeakBagTicket -> IO ()
remove (WeakBagTicket Weak a
w a
_) = {-# SCC "remove" #-} Weak a -> IO ()
forall v. Weak v -> IO ()
finalize Weak a
w
--TODO: Should 'remove' also drop the reference to the item?

--TODO: can/should we provide a null WeakBagTicket?