-- | Thin wrappers over @io-classes@ that attach a label to the STM
-- primitive or thread they create.
--
-- The labels surface in @io-sim@ traces and as GHC thread labels, which
-- makes deadlock and liveness diagnosis tractable. Applying them everywhere
-- a 'TVar'\/'TMVar'\/'TQueue'\/'TBQueue' is allocated or a thread is spawned
-- is a discipline worth keeping; these helpers make it the path of least
-- resistance.
module Control.Concurrent.Class.Labelled (
  -- * Labelling STM variables
  newLabelledTVar,
  newLabelledTVarIO,
  newLabelledEmptyTMVar,
  newLabelledEmptyTMVarIO,
  newLabelledTQueue,
  newLabelledTQueueIO,
  newLabelledTBQueue,
  newLabelledTBQueueIO,

  -- * Labelling threads
  asyncLabelled,
  raceLabelled,
  raceLabelled_,
  withAsyncLabelled,
  concurrentlyLabelled,
  concurrentlyLabelled_,
) where

import Control.Concurrent.Class.MonadSTM (
  MonadLabelledSTM,
  STM,
  TBQueue,
  TMVar,
  TQueue,
  TVar,
  atomically,
  labelTBQueue,
  labelTMVar,
  labelTQueue,
  labelTVar,
  newEmptyTMVar,
  newTBQueue,
  newTQueue,
  newTVar,
 )
import Control.Monad (void)
import Control.Monad.Class.MonadAsync (Async, MonadAsync, async, concurrently, race, withAsync)
import Control.Monad.Class.MonadFork (labelThisThread)
import Numeric.Natural (Natural)

-- * Helpers for labeling TVar

newLabelledTVar :: MonadLabelledSTM m => String -> a -> STM m (TVar m a)
newLabelledTVar :: forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> a -> STM m (TVar m a)
newLabelledTVar String
lbl a
val = do
  TVar m a
tv <- a -> STM m (TVar m a)
forall a. a -> STM m (TVar m a)
forall (m :: * -> *) a. MonadSTM m => a -> STM m (TVar m a)
newTVar a
val
  TVar m a -> String -> STM m ()
forall a. TVar m a -> String -> STM m ()
forall (m :: * -> *) a.
MonadLabelledSTM m =>
TVar m a -> String -> STM m ()
labelTVar TVar m a
tv String
lbl
  TVar m a -> STM m (TVar m a)
forall a. a -> STM m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TVar m a
tv

newLabelledTVarIO :: MonadLabelledSTM m => String -> a -> m (TVar m a)
newLabelledTVarIO :: forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> a -> m (TVar m a)
newLabelledTVarIO = (STM m (TVar m a) -> m (TVar m a)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically .) ((a -> STM m (TVar m a)) -> a -> m (TVar m a))
-> (String -> a -> STM m (TVar m a)) -> String -> a -> m (TVar m a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> a -> STM m (TVar m a)
forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> a -> STM m (TVar m a)
newLabelledTVar

-- * Helpers for labeling TMVar

newLabelledEmptyTMVar :: MonadLabelledSTM m => String -> STM m (TMVar m a)
newLabelledEmptyTMVar :: forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> STM m (TMVar m a)
newLabelledEmptyTMVar String
lbl = do
  TMVar m a
tmv <- STM m (TMVar m a)
forall a. STM m (TMVar m a)
forall (m :: * -> *) a. MonadSTM m => STM m (TMVar m a)
newEmptyTMVar
  TMVar m a -> String -> STM m ()
forall a. TMVar m a -> String -> STM m ()
forall (m :: * -> *) a.
MonadLabelledSTM m =>
TMVar m a -> String -> STM m ()
labelTMVar TMVar m a
tmv String
lbl
  TMVar m a -> STM m (TMVar m a)
forall a. a -> STM m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TMVar m a
tmv

newLabelledEmptyTMVarIO :: MonadLabelledSTM m => String -> m (TMVar m a)
newLabelledEmptyTMVarIO :: forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> m (TMVar m a)
newLabelledEmptyTMVarIO = STM m (TMVar m a) -> m (TMVar m a)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m (TMVar m a) -> m (TMVar m a))
-> (String -> STM m (TMVar m a)) -> String -> m (TMVar m a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> STM m (TMVar m a)
forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> STM m (TMVar m a)
newLabelledEmptyTMVar

-- * Helpers for labeling TQueue

newLabelledTQueue :: MonadLabelledSTM m => String -> STM m (TQueue m a)
newLabelledTQueue :: forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> STM m (TQueue m a)
newLabelledTQueue String
lbl = do
  TQueue m a
q <- STM m (TQueue m a)
forall a. STM m (TQueue m a)
forall (m :: * -> *) a. MonadSTM m => STM m (TQueue m a)
newTQueue
  TQueue m a -> String -> STM m ()
forall a. TQueue m a -> String -> STM m ()
forall (m :: * -> *) a.
MonadLabelledSTM m =>
TQueue m a -> String -> STM m ()
labelTQueue TQueue m a
q String
lbl
  TQueue m a -> STM m (TQueue m a)
forall a. a -> STM m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TQueue m a
q

newLabelledTQueueIO :: MonadLabelledSTM m => String -> m (TQueue m a)
newLabelledTQueueIO :: forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> m (TQueue m a)
newLabelledTQueueIO = STM m (TQueue m a) -> m (TQueue m a)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m (TQueue m a) -> m (TQueue m a))
-> (String -> STM m (TQueue m a)) -> String -> m (TQueue m a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> STM m (TQueue m a)
forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> STM m (TQueue m a)
newLabelledTQueue

-- * Helpers for labeling TBQueue

newLabelledTBQueue :: MonadLabelledSTM m => String -> Natural -> STM m (TBQueue m a)
newLabelledTBQueue :: forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> Natural -> STM m (TBQueue m a)
newLabelledTBQueue String
lbl Natural
capacity = do
  TBQueue m a
bq <- Natural -> STM m (TBQueue m a)
forall a. Natural -> STM m (TBQueue m a)
forall (m :: * -> *) a.
MonadSTM m =>
Natural -> STM m (TBQueue m a)
newTBQueue Natural
capacity
  TBQueue m a -> String -> STM m ()
forall a. TBQueue m a -> String -> STM m ()
forall (m :: * -> *) a.
MonadLabelledSTM m =>
TBQueue m a -> String -> STM m ()
labelTBQueue TBQueue m a
bq String
lbl
  TBQueue m a -> STM m (TBQueue m a)
forall a. a -> STM m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TBQueue m a
bq

newLabelledTBQueueIO :: MonadLabelledSTM m => String -> Natural -> m (TBQueue m a)
newLabelledTBQueueIO :: forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> Natural -> m (TBQueue m a)
newLabelledTBQueueIO = (STM m (TBQueue m a) -> m (TBQueue m a)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically .) ((Natural -> STM m (TBQueue m a)) -> Natural -> m (TBQueue m a))
-> (String -> Natural -> STM m (TBQueue m a))
-> String
-> Natural
-> m (TBQueue m a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Natural -> STM m (TBQueue m a)
forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> Natural -> STM m (TBQueue m a)
newLabelledTBQueue

-- * Helpers for labeling Threads

raceLabelled :: MonadAsync m => (String, m a) -> (String, m b) -> m (Either a b)
raceLabelled :: forall (m :: * -> *) a b.
MonadAsync m =>
(String, m a) -> (String, m b) -> m (Either a b)
raceLabelled (String
lblA, m a
mA) (String
lblB, m b
mB) =
  m a -> m b -> m (Either a b)
forall a b. m a -> m b -> m (Either a b)
forall (m :: * -> *) a b.
MonadAsync m =>
m a -> m b -> m (Either a b)
race
    (String -> m ()
forall (m :: * -> *). MonadThread m => String -> m ()
labelThisThread String
lblA m () -> m a -> m a
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m a
mA)
    (String -> m ()
forall (m :: * -> *). MonadThread m => String -> m ()
labelThisThread String
lblB m () -> m b -> m b
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m b
mB)

raceLabelled_ :: MonadAsync m => (String, m a) -> (String, m b) -> m ()
raceLabelled_ :: forall (m :: * -> *) a b.
MonadAsync m =>
(String, m a) -> (String, m b) -> m ()
raceLabelled_ = (m (Either a b) -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void .) (((String, m b) -> m (Either a b)) -> (String, m b) -> m ())
-> ((String, m a) -> (String, m b) -> m (Either a b))
-> (String, m a)
-> (String, m b)
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String, m a) -> (String, m b) -> m (Either a b)
forall (m :: * -> *) a b.
MonadAsync m =>
(String, m a) -> (String, m b) -> m (Either a b)
raceLabelled

withAsyncLabelled :: MonadAsync m => (String, m a) -> (Async m a -> m b) -> m b
withAsyncLabelled :: forall (m :: * -> *) a b.
MonadAsync m =>
(String, m a) -> (Async m a -> m b) -> m b
withAsyncLabelled (String
lbl, m a
ma) = m a -> (Async m a -> m b) -> m b
forall a b. m a -> (Async m a -> m b) -> m b
forall (m :: * -> *) a b.
MonadAsync m =>
m a -> (Async m a -> m b) -> m b
withAsync (String -> m ()
forall (m :: * -> *). MonadThread m => String -> m ()
labelThisThread String
lbl m () -> m a -> m a
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m a
ma)

concurrentlyLabelled :: MonadAsync m => (String, m a) -> (String, m b) -> m (a, b)
concurrentlyLabelled :: forall (m :: * -> *) a b.
MonadAsync m =>
(String, m a) -> (String, m b) -> m (a, b)
concurrentlyLabelled (String
lblA, m a
mA) (String
lblB, m b
mB) =
  m a -> m b -> m (a, b)
forall a b. m a -> m b -> m (a, b)
forall (m :: * -> *) a b. MonadAsync m => m a -> m b -> m (a, b)
concurrently
    (String -> m ()
forall (m :: * -> *). MonadThread m => String -> m ()
labelThisThread String
lblA m () -> m a -> m a
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m a
mA)
    (String -> m ()
forall (m :: * -> *). MonadThread m => String -> m ()
labelThisThread String
lblB m () -> m b -> m b
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m b
mB)

concurrentlyLabelled_ :: MonadAsync m => (String, m a) -> (String, m b) -> m ()
concurrentlyLabelled_ :: forall (m :: * -> *) a b.
MonadAsync m =>
(String, m a) -> (String, m b) -> m ()
concurrentlyLabelled_ = (m (a, b) -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void .) (((String, m b) -> m (a, b)) -> (String, m b) -> m ())
-> ((String, m a) -> (String, m b) -> m (a, b))
-> (String, m a)
-> (String, m b)
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String, m a) -> (String, m b) -> m (a, b)
forall (m :: * -> *) a b.
MonadAsync m =>
(String, m a) -> (String, m b) -> m (a, b)
concurrentlyLabelled

asyncLabelled :: MonadAsync m => String -> m a -> m (Async m a)
asyncLabelled :: forall (m :: * -> *) a.
MonadAsync m =>
String -> m a -> m (Async m a)
asyncLabelled String
lbl m a
mA = m a -> m (Async m a)
forall a. m a -> m (Async m a)
forall (m :: * -> *) a. MonadAsync m => m a -> m (Async m a)
async (m a -> m (Async m a)) -> m a -> m (Async m a)
forall a b. (a -> b) -> a -> b
$ String -> m ()
forall (m :: * -> *). MonadThread m => String -> m ()
labelThisThread String
lbl m () -> m a -> m a
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m a
mA