{-# LANGUAGE DeriveAnyClass #-}

-- | @LVar@ is like @Control.Concurrent.STM.TMVar@ but with a capability for
-- listening to its changes.
module Data.LVar
  ( -- * Types
    LVar,
    ListenerId,

    -- * Creating a LVar
    new,
    empty,

    -- * Modifying a LVar
    get,
    set,
    modify,

    -- * Listening to a LVar
    addListener,
    listenNext,
    removeListener,
  )
where

import Control.Exception (throw)
import qualified Data.Map.Strict as Map
import Prelude hiding (empty, get, modify)

-- A mutable variable (like @TMVar@), changes to which can be listened to from
-- multiple threads.
data LVar a = LVar
  { -- | A value that changes over time
    LVar a -> TMVar a
lvarCurrent :: TMVar a,
    -- | Subscribers listening on changes to the value
    LVar a -> TMVar (Map ListenerId (TMVar ()))
lvarListeners :: TMVar (Map ListenerId (TMVar ()))
  }

type ListenerId = Int

-- | Create a new @LVar@ with the given initial value
new :: forall a m. MonadIO m => a -> m (LVar a)
new :: a -> m (LVar a)
new a
val = do
  TMVar a -> TMVar (Map ListenerId (TMVar ())) -> LVar a
forall a. TMVar a -> TMVar (Map ListenerId (TMVar ())) -> LVar a
LVar (TMVar a -> TMVar (Map ListenerId (TMVar ())) -> LVar a)
-> m (TMVar a) -> m (TMVar (Map ListenerId (TMVar ())) -> LVar a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m (TMVar a)
forall (m :: * -> *) a. MonadIO m => a -> m (TMVar a)
newTMVarIO a
val m (TMVar (Map ListenerId (TMVar ())) -> LVar a)
-> m (TMVar (Map ListenerId (TMVar ()))) -> m (LVar a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Map ListenerId (TMVar ()) -> m (TMVar (Map ListenerId (TMVar ())))
forall (m :: * -> *) a. MonadIO m => a -> m (TMVar a)
newTMVarIO Map ListenerId (TMVar ())
forall a. Monoid a => a
mempty

-- | Like @new@, but there is no initial value. A @get@ will block until an
-- initial value is set using @set@ or @modify@
empty :: MonadIO m => m (LVar a)
empty :: m (LVar a)
empty =
  TMVar a -> TMVar (Map ListenerId (TMVar ())) -> LVar a
forall a. TMVar a -> TMVar (Map ListenerId (TMVar ())) -> LVar a
LVar (TMVar a -> TMVar (Map ListenerId (TMVar ())) -> LVar a)
-> m (TMVar a) -> m (TMVar (Map ListenerId (TMVar ())) -> LVar a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (TMVar a)
forall (m :: * -> *) a. MonadIO m => m (TMVar a)
newEmptyTMVarIO m (TMVar (Map ListenerId (TMVar ())) -> LVar a)
-> m (TMVar (Map ListenerId (TMVar ()))) -> m (LVar a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Map ListenerId (TMVar ()) -> m (TMVar (Map ListenerId (TMVar ())))
forall (m :: * -> *) a. MonadIO m => a -> m (TMVar a)
newTMVarIO Map ListenerId (TMVar ())
forall a. Monoid a => a
mempty

-- | Get the value of the @LVar@
get :: MonadIO m => LVar a -> m a
get :: LVar a -> m a
get LVar a
v =
  STM a -> m a
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM a -> m a) -> STM a -> m a
forall a b. (a -> b) -> a -> b
$ TMVar a -> STM a
forall a. TMVar a -> STM a
readTMVar (TMVar a -> STM a) -> TMVar a -> STM a
forall a b. (a -> b) -> a -> b
$ LVar a -> TMVar a
forall a. LVar a -> TMVar a
lvarCurrent LVar a
v

-- | Set the @LVar@ value; active listeners are automatically notifed.
set :: MonadIO m => LVar a -> a -> m ()
set :: LVar a -> a -> m ()
set LVar a
v a
val = do
  STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> m ()) -> STM () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    let var :: TMVar a
var = LVar a -> TMVar a
forall a. LVar a -> TMVar a
lvarCurrent LVar a
v
    TMVar a -> STM Bool
forall a. TMVar a -> STM Bool
isEmptyTMVar TMVar a
var STM Bool -> (Bool -> STM ()) -> STM ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Bool
True -> TMVar a -> a -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar TMVar a
var a
val
      Bool
False -> STM a -> STM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (STM a -> STM ()) -> STM a -> STM ()
forall a b. (a -> b) -> a -> b
$ TMVar a -> a -> STM a
forall a. TMVar a -> a -> STM a
swapTMVar TMVar a
var a
val
    LVar a -> STM ()
forall a. LVar a -> STM ()
notifyListeners LVar a
v

-- | Modify the @LVar@ value; active listeners are automatically notified.
modify :: MonadIO m => LVar a -> (a -> a) -> m ()
modify :: LVar a -> (a -> a) -> m ()
modify LVar a
v a -> a
f = do
  STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> m ()) -> STM () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    a
curr <- TMVar a -> STM a
forall a. TMVar a -> STM a
readTMVar (LVar a -> TMVar a
forall a. LVar a -> TMVar a
lvarCurrent LVar a
v)
    STM a -> STM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (STM a -> STM ()) -> STM a -> STM ()
forall a b. (a -> b) -> a -> b
$ TMVar a -> a -> STM a
forall a. TMVar a -> a -> STM a
swapTMVar (LVar a -> TMVar a
forall a. LVar a -> TMVar a
lvarCurrent LVar a
v) (a -> a
f a
curr)
    LVar a -> STM ()
forall a. LVar a -> STM ()
notifyListeners LVar a
v

notifyListeners :: LVar a -> STM ()
notifyListeners :: LVar a -> STM ()
notifyListeners LVar a
v' = do
  Map ListenerId (TMVar ())
subs <- TMVar (Map ListenerId (TMVar ()))
-> STM (Map ListenerId (TMVar ()))
forall a. TMVar a -> STM a
readTMVar (TMVar (Map ListenerId (TMVar ()))
 -> STM (Map ListenerId (TMVar ())))
-> TMVar (Map ListenerId (TMVar ()))
-> STM (Map ListenerId (TMVar ()))
forall a b. (a -> b) -> a -> b
$ LVar a -> TMVar (Map ListenerId (TMVar ()))
forall a. LVar a -> TMVar (Map ListenerId (TMVar ()))
lvarListeners LVar a
v'
  [TMVar ()] -> (TMVar () -> STM Bool) -> STM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Map ListenerId (TMVar ()) -> [TMVar ()]
forall k a. Map k a -> [a]
Map.elems Map ListenerId (TMVar ())
subs) ((TMVar () -> STM Bool) -> STM ())
-> (TMVar () -> STM Bool) -> STM ()
forall a b. (a -> b) -> a -> b
$ \TMVar ()
subVar -> do
    TMVar () -> () -> STM Bool
forall a. TMVar a -> a -> STM Bool
tryPutTMVar TMVar ()
subVar ()

data ListenerDead = ListenerDead
  deriving (Show ListenerDead
Typeable ListenerDead
Typeable ListenerDead
-> Show ListenerDead
-> (ListenerDead -> SomeException)
-> (SomeException -> Maybe ListenerDead)
-> (ListenerDead -> String)
-> Exception ListenerDead
SomeException -> Maybe ListenerDead
ListenerDead -> String
ListenerDead -> SomeException
forall e.
Typeable e
-> Show e
-> (e -> SomeException)
-> (SomeException -> Maybe e)
-> (e -> String)
-> Exception e
displayException :: ListenerDead -> String
$cdisplayException :: ListenerDead -> String
fromException :: SomeException -> Maybe ListenerDead
$cfromException :: SomeException -> Maybe ListenerDead
toException :: ListenerDead -> SomeException
$ctoException :: ListenerDead -> SomeException
$cp2Exception :: Show ListenerDead
$cp1Exception :: Typeable ListenerDead
Exception, ListenerId -> ListenerDead -> ShowS
[ListenerDead] -> ShowS
ListenerDead -> String
(ListenerId -> ListenerDead -> ShowS)
-> (ListenerDead -> String)
-> ([ListenerDead] -> ShowS)
-> Show ListenerDead
forall a.
(ListenerId -> a -> ShowS)
-> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ListenerDead] -> ShowS
$cshowList :: [ListenerDead] -> ShowS
show :: ListenerDead -> String
$cshow :: ListenerDead -> String
showsPrec :: ListenerId -> ListenerDead -> ShowS
$cshowsPrec :: ListenerId -> ListenerDead -> ShowS
Show)

-- | Create a listener for changes to the @LVar@, as they are set by @set@ or
-- @modify@ from this time onwards.
--
-- You must call @listenNext@ to get the next updated value (or current value if
-- there is one).
--
-- Returns a @ListenerId@ that can be used to stop listening later (via
-- @removeListener@)
addListener ::
  MonadIO m =>
  LVar a ->
  m ListenerId
addListener :: LVar a -> m ListenerId
addListener LVar a
v = do
  STM ListenerId -> m ListenerId
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM ListenerId -> m ListenerId) -> STM ListenerId -> m ListenerId
forall a b. (a -> b) -> a -> b
$ do
    Map ListenerId (TMVar ())
subs <- TMVar (Map ListenerId (TMVar ()))
-> STM (Map ListenerId (TMVar ()))
forall a. TMVar a -> STM a
readTMVar (TMVar (Map ListenerId (TMVar ()))
 -> STM (Map ListenerId (TMVar ())))
-> TMVar (Map ListenerId (TMVar ()))
-> STM (Map ListenerId (TMVar ()))
forall a b. (a -> b) -> a -> b
$ LVar a -> TMVar (Map ListenerId (TMVar ()))
forall a. LVar a -> TMVar (Map ListenerId (TMVar ()))
lvarListeners LVar a
v
    let nextIdx :: ListenerId
nextIdx = ListenerId
-> ((ListenerId, TMVar ()) -> ListenerId)
-> Maybe (ListenerId, TMVar ())
-> ListenerId
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ListenerId
1 (ListenerId -> ListenerId
forall a. Enum a => a -> a
succ (ListenerId -> ListenerId)
-> ((ListenerId, TMVar ()) -> ListenerId)
-> (ListenerId, TMVar ())
-> ListenerId
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ListenerId, TMVar ()) -> ListenerId
forall a b. (a, b) -> a
fst) (Maybe (ListenerId, TMVar ()) -> ListenerId)
-> Maybe (ListenerId, TMVar ()) -> ListenerId
forall a b. (a -> b) -> a -> b
$ Map ListenerId (TMVar ()) -> Maybe (ListenerId, TMVar ())
forall k a. Map k a -> Maybe (k, a)
Map.lookupMax Map ListenerId (TMVar ())
subs
    TMVar ()
notify <-
      TMVar a -> STM (Maybe a)
forall a. TMVar a -> STM (Maybe a)
tryReadTMVar (LVar a -> TMVar a
forall a. LVar a -> TMVar a
lvarCurrent LVar a
v) STM (Maybe a) -> (Maybe a -> STM (TMVar ())) -> STM (TMVar ())
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Maybe a
Nothing -> STM (TMVar ())
forall a. STM (TMVar a)
newEmptyTMVar
        -- As a value is already available, send that as first notification.
        --
        -- NOTE: Creating a TMVar that is "full" ensures that we send a current
        -- (which is not empty) value on @listenNext@).
        Just a
_ -> () -> STM (TMVar ())
forall a. a -> STM (TMVar a)
newTMVar ()
    STM (Map ListenerId (TMVar ())) -> STM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (STM (Map ListenerId (TMVar ())) -> STM ())
-> STM (Map ListenerId (TMVar ())) -> STM ()
forall a b. (a -> b) -> a -> b
$ TMVar (Map ListenerId (TMVar ()))
-> Map ListenerId (TMVar ()) -> STM (Map ListenerId (TMVar ()))
forall a. TMVar a -> a -> STM a
swapTMVar (LVar a -> TMVar (Map ListenerId (TMVar ()))
forall a. LVar a -> TMVar (Map ListenerId (TMVar ()))
lvarListeners LVar a
v) (Map ListenerId (TMVar ()) -> STM (Map ListenerId (TMVar ())))
-> Map ListenerId (TMVar ()) -> STM (Map ListenerId (TMVar ()))
forall a b. (a -> b) -> a -> b
$ ListenerId
-> TMVar ()
-> Map ListenerId (TMVar ())
-> Map ListenerId (TMVar ())
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert ListenerId
nextIdx TMVar ()
notify Map ListenerId (TMVar ())
subs
    ListenerId -> STM ListenerId
forall (f :: * -> *) a. Applicative f => a -> f a
pure ListenerId
nextIdx

-- | Listen for the next value update (since the last @listenNext@ or
-- @addListener@). Unless the @LVar@ was empty when @addListener@ was invoked,
-- the first invocation of @listenNext@ will return the current value even if
-- there wasn't an update.  Therefore, the *first* call to @listenNext@ will
-- *always* return immediately, unless the @LVar@ is empty.
--
-- Call this in a loop to listen on a series of updates.
--
-- Throws @ListenerDead@ if called with a @ListenerId@ that got already removed
-- by @removeListener@.
listenNext :: MonadIO m => LVar a -> ListenerId -> m a
listenNext :: LVar a -> ListenerId -> m a
listenNext LVar a
v ListenerId
idx = do
  STM a -> m a
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM a -> m a) -> STM a -> m a
forall a b. (a -> b) -> a -> b
$ do
    LVar a -> ListenerId -> STM (Maybe (TMVar ()))
forall a. LVar a -> ListenerId -> STM (Maybe (TMVar ()))
lookupListener LVar a
v ListenerId
idx STM (Maybe (TMVar ())) -> (Maybe (TMVar ()) -> STM a) -> STM a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Maybe (TMVar ())
Nothing ->
        -- FIXME: can we avoid this by design?
        ListenerDead -> STM a
forall a e. Exception e => e -> a
throw ListenerDead
ListenerDead
      Just TMVar ()
listenVar -> do
        TMVar () -> STM ()
forall a. TMVar a -> STM a
takeTMVar TMVar ()
listenVar
        TMVar a -> STM a
forall a. TMVar a -> STM a
readTMVar (LVar a -> TMVar a
forall a. LVar a -> TMVar a
lvarCurrent LVar a
v)
  where
    lookupListener :: LVar a -> ListenerId -> STM (Maybe (TMVar ()))
    lookupListener :: LVar a -> ListenerId -> STM (Maybe (TMVar ()))
lookupListener LVar a
v' ListenerId
lId = do
      ListenerId -> Map ListenerId (TMVar ()) -> Maybe (TMVar ())
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup ListenerId
lId (Map ListenerId (TMVar ()) -> Maybe (TMVar ()))
-> STM (Map ListenerId (TMVar ())) -> STM (Maybe (TMVar ()))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TMVar (Map ListenerId (TMVar ()))
-> STM (Map ListenerId (TMVar ()))
forall a. TMVar a -> STM a
readTMVar (LVar a -> TMVar (Map ListenerId (TMVar ()))
forall a. LVar a -> TMVar (Map ListenerId (TMVar ()))
lvarListeners LVar a
v')

-- | Stop listening to the @LVar@
removeListener :: MonadIO m => LVar a -> ListenerId -> m ()
removeListener :: LVar a -> ListenerId -> m ()
removeListener LVar a
v ListenerId
lId = do
  STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> m ()) -> STM () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    Map ListenerId (TMVar ())
subs <- TMVar (Map ListenerId (TMVar ()))
-> STM (Map ListenerId (TMVar ()))
forall a. TMVar a -> STM a
readTMVar (TMVar (Map ListenerId (TMVar ()))
 -> STM (Map ListenerId (TMVar ())))
-> TMVar (Map ListenerId (TMVar ()))
-> STM (Map ListenerId (TMVar ()))
forall a b. (a -> b) -> a -> b
$ LVar a -> TMVar (Map ListenerId (TMVar ()))
forall a. LVar a -> TMVar (Map ListenerId (TMVar ()))
lvarListeners LVar a
v
    Maybe (TMVar ()) -> (TMVar () -> STM ()) -> STM ()
forall (f :: * -> *) a.
Applicative f =>
Maybe a -> (a -> f ()) -> f ()
whenJust (ListenerId -> Map ListenerId (TMVar ()) -> Maybe (TMVar ())
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup ListenerId
lId Map ListenerId (TMVar ())
subs) ((TMVar () -> STM ()) -> STM ()) -> (TMVar () -> STM ()) -> STM ()
forall a b. (a -> b) -> a -> b
$ \TMVar ()
_sub -> do
      STM (Map ListenerId (TMVar ())) -> STM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (STM (Map ListenerId (TMVar ())) -> STM ())
-> STM (Map ListenerId (TMVar ())) -> STM ()
forall a b. (a -> b) -> a -> b
$ TMVar (Map ListenerId (TMVar ()))
-> Map ListenerId (TMVar ()) -> STM (Map ListenerId (TMVar ()))
forall a. TMVar a -> a -> STM a
swapTMVar (LVar a -> TMVar (Map ListenerId (TMVar ()))
forall a. LVar a -> TMVar (Map ListenerId (TMVar ()))
lvarListeners LVar a
v) (Map ListenerId (TMVar ()) -> STM (Map ListenerId (TMVar ())))
-> Map ListenerId (TMVar ()) -> STM (Map ListenerId (TMVar ()))
forall a b. (a -> b) -> a -> b
$ ListenerId
-> Map ListenerId (TMVar ()) -> Map ListenerId (TMVar ())
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete ListenerId
lId Map ListenerId (TMVar ())
subs