{-# 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
  { WeakBag a -> IORef Int
_weakBag_nextId :: {-# UNPACK #-} !(IORef Int) --TODO: what if this wraps around?
  , 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 :: a
-> WeakBag a -> IORef (Weak b) -> (b -> IO ()) -> IO WeakBagTicket
insert a :: a
a (WeakBag nextId :: IORef Int
nextId children :: IORef (IntMap (Weak a))
children) wbRef :: IORef (Weak b)
wbRef finalizer :: b -> IO ()
finalizer = {-# SCC "insert" #-} do
  a
a' <- a -> IO a
forall a. a -> IO a
evaluate a
a
  IORef (Weak b)
wbRef' <- IORef (Weak b) -> IO (IORef (Weak b))
forall a. a -> IO a
evaluate IORef (Weak b)
wbRef
  Int
myId <- IORef Int -> (Int -> (Int, Int)) -> IO Int
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Int
nextId ((Int -> (Int, Int)) -> IO Int) -> (Int -> (Int, Int)) -> IO Int
forall a b. (a -> b) -> a -> b
$ \n :: Int
n -> (Int -> Int
forall a. Enum a => a -> a
succ Int
n, Int
n)
  let cleanup :: IO ()
cleanup = do
        Weak b
wb <- IORef (Weak b) -> IO (Weak b)
forall a. IORef a -> IO a
readIORef IORef (Weak b)
wbRef'
        Maybe b
mb <- Weak b -> IO (Maybe b)
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak b
wb
        Maybe b -> (b -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe b
mb ((b -> IO ()) -> IO ()) -> (b -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \b :: b
b -> do
          IntMap (Weak a)
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
$ \cs :: 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)
          Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (IntMap (Weak a) -> Bool
forall a. IntMap a -> Bool
IntMap.null IntMap (Weak a)
csWithoutMe) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ b -> IO ()
finalizer b
b
  Weak a
wa <- a -> Maybe (IO ()) -> IO (Weak a)
forall k. k -> Maybe (IO ()) -> IO (Weak k)
mkWeakPtr a
a' (Maybe (IO ()) -> IO (Weak a)) -> Maybe (IO ()) -> IO (Weak a)
forall a b. (a -> b) -> a -> b
$ IO () -> Maybe (IO ())
forall a. a -> Maybe a
Just IO ()
cleanup
  IORef (IntMap (Weak a))
-> (IntMap (Weak a) -> (IntMap (Weak a), ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (IntMap (Weak a))
children ((IntMap (Weak a) -> (IntMap (Weak a), ())) -> IO ())
-> (IntMap (Weak a) -> (IntMap (Weak a), ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \cs :: 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, ())
  WeakBagTicket -> IO WeakBagTicket
forall (m :: * -> *) a. Monad m => a -> m a
return (WeakBagTicket -> IO WeakBagTicket)
-> WeakBagTicket -> IO WeakBagTicket
forall a b. (a -> b) -> a -> b
$ $WWeakBagTicket :: forall a. Weak a -> a -> WeakBagTicket
WeakBagTicket
    { _weakBagTicket_weakItem :: Weak a
_weakBagTicket_weakItem = Weak a
wa
    , _weakBagTicket_item :: a
_weakBagTicket_item = a
a'
    }

-- | Create an empty 'WeakBag'.
{-# INLINE empty #-}
empty :: IO (WeakBag a)
empty :: IO (WeakBag a)
empty = {-# SCC "empty" #-} do
  IORef Int
nextId <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef 1
  IORef (IntMap (Weak a))
children <- IntMap (Weak a) -> IO (IORef (IntMap (Weak a)))
forall a. a -> IO (IORef a)
newIORef IntMap (Weak a)
forall a. IntMap a
IntMap.empty
  let bag :: WeakBag a
bag = $WWeakBag :: forall a. IORef Int -> IORef (IntMap (Weak a)) -> WeakBag a
WeakBag
        { _weakBag_nextId :: IORef Int
_weakBag_nextId = IORef Int
nextId
        , _weakBag_children :: IORef (IntMap (Weak a))
_weakBag_children = IORef (IntMap (Weak a))
children
        }
  WeakBag a -> IO (WeakBag a)
forall (m :: * -> *) a. Monad m => a -> m a
return WeakBag a
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 :: a
-> IORef (Weak b) -> (b -> IO ()) -> IO (WeakBag a, WeakBagTicket)
singleton a :: a
a wbRef :: IORef (Weak b)
wbRef finalizer :: b -> IO ()
finalizer = {-# SCC "singleton" #-} do
  WeakBag a
bag <- IO (WeakBag a)
forall a. IO (WeakBag a)
empty
  WeakBagTicket
ticket <- a
-> WeakBag a -> IORef (Weak b) -> (b -> IO ()) -> IO WeakBagTicket
forall a b.
a
-> WeakBag a -> IORef (Weak b) -> (b -> IO ()) -> IO WeakBagTicket
insert a
a WeakBag a
bag IORef (Weak b)
wbRef b -> IO ()
finalizer
  (WeakBag a, WeakBagTicket) -> IO (WeakBag a, WeakBagTicket)
forall (m :: * -> *) a. Monad m => a -> m a
return (WeakBag a
bag, WeakBagTicket
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_ :: WeakBag a -> (a -> m ()) -> m ()
traverse_ (WeakBag _ children :: IORef (IntMap (Weak a))
children) f :: a -> m ()
f = {-# SCC "traverse" #-} do
  IntMap (Weak a)
cs <- IO (IntMap (Weak a)) -> m (IntMap (Weak 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
  IntMap (Weak a) -> (Weak a -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ IntMap (Weak a)
cs ((Weak a -> m ()) -> m ()) -> (Weak a -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \c :: Weak a
c -> do
    Maybe a
ma <- IO (Maybe a) -> m (Maybe 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
    (a -> m ()) -> Maybe a -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ a -> m ()
f Maybe a
ma

{-# DEPRECATED traverse "Use 'traverse_' instead" #-}
traverse :: MonadIO m => WeakBag a -> (a -> m ()) -> m ()
traverse :: 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 w :: Weak a
w _) = {-# 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?