    {-# LANGUAGE NoImplicitPrelude #-}
    {-# LANGUAGE RankNTypes        #-}

    -- | Companion threads, such as for printing messages saying we're still busy.

    module Control.Concurrent.Companion
      ( withCompanion
      , onCompanionDone
      , Companion
      , Delay
      , StopCompanion
      ) where

    import           RIO

    -- | A companion thread which can perform arbitrary actions as well as delay.

    type Companion m = Delay -> m ()

    -- | Delay the given number of microseconds. If 'StopCompanion' is triggered

    -- before the timer completes, a 'CompanionDone' exception will be thrown (which

    -- is caught internally by 'withCompanion').

    type Delay = forall mio. MonadIO mio => Int -> mio ()

    -- | Tell the 'Companion' to stop. The next time 'Delay' is called, or if a

    -- 'Delay' is currently blocking, the 'Companion' thread will exit with a

    -- 'CompanionDone' exception.

    type StopCompanion m = m ()

    -- | When a delay was interrupted because we're told to stop, perform this

    -- action.

    onCompanionDone
      :: MonadUnliftIO m
      => m () -- ^ the delay

      -> m () -- ^ action to perform

      -> m ()
    onCompanionDone :: forall (m :: * -> *). MonadUnliftIO m => m () -> m () -> m ()
    onCompanionDone m ()
    theDelay m ()
    theAction =
      m ()
    theDelay m () -> (CompanionDone -> m ()) -> m ()
    forall (m :: * -> *) e a b.
    (MonadUnliftIO m, Exception e) =>
    m a -> (e -> m b) -> m a
    `withException` \CompanionDone
    CompanionDone -> m ()
    theAction

    -- | Internal exception used by 'withCompanion' to allow short-circuiting of the

    -- 'Companion'. Should not be used outside of this module.

    data CompanionDone = CompanionDone
      deriving (Int -> CompanionDone -> ShowS
    [CompanionDone] -> ShowS
    CompanionDone -> String
    (Int -> CompanionDone -> ShowS)
    -> (CompanionDone -> String)
    -> ([CompanionDone] -> ShowS)
    -> Show CompanionDone
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> CompanionDone -> ShowS
    showsPrec :: Int -> CompanionDone -> ShowS
    $cshow :: CompanionDone -> String
    show :: CompanionDone -> String
    $cshowList :: [CompanionDone] -> ShowS
    showList :: [CompanionDone] -> ShowS
    Show, Typeable)

    instance Exception CompanionDone

    -- | Keep running the 'Companion' action until either the inner action completes

    -- or calls the 'StopCompanion' action. This can be used to give the user status

    -- information while running a long running operations.

    withCompanion ::
         forall m a. MonadUnliftIO m
      => Companion m
      -> (StopCompanion m -> m a)
      -> m a
    withCompanion :: forall (m :: * -> *) a.
    MonadUnliftIO m =>
    Companion m -> (StopCompanion m -> m a) -> m a
    withCompanion Companion m
    companion StopCompanion m -> m a
    inner = do
      -- Variable to indicate 'Delay'ing should result in a 'CompanionDone'

      -- exception.

      TVar Bool
    shouldStopVar <- Bool -> m (TVar Bool)
    forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
    newTVarIO Bool
    False
      let -- Relatively simple: set shouldStopVar to True

          stopCompanion :: StopCompanion m
    stopCompanion = STM () -> StopCompanion m
    forall (m :: * -> *) a. MonadIO m => STM a -> m a
    atomically (STM () -> StopCompanion m) -> STM () -> StopCompanion m
    forall a b. (a -> b) -> a -> b
    $ TVar Bool -> Bool -> STM ()
    forall a. TVar a -> a -> STM ()
    writeTVar TVar Bool
    shouldStopVar Bool
    True

          delay :: Delay
          delay :: Delay
    delay Int
    usec = do
            -- Register a delay with the runtime system

            TVar Bool
    delayDoneVar <- Int -> mio (TVar Bool)
    forall (m :: * -> *). MonadIO m => Int -> m (TVar Bool)
    registerDelay Int
    usec
            mio (mio ()) -> mio ()
    forall (m :: * -> *) a. Monad m => m (m a) -> m a
    join (mio (mio ()) -> mio ()) -> mio (mio ()) -> mio ()
    forall a b. (a -> b) -> a -> b
    $ STM (mio ()) -> mio (mio ())
    forall (m :: * -> *) a. MonadIO m => STM a -> m a
    atomically (STM (mio ()) -> mio (mio ())) -> STM (mio ()) -> mio (mio ())
    forall a b. (a -> b) -> a -> b
    $
              -- Delay has triggered, keep going

              (() -> mio ()
    forall a. a -> mio a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure () mio () -> STM () -> STM (mio ())
    forall a b. a -> STM b -> STM a
    forall (f :: * -> *) a b. Functor f => a -> f b -> f a
    <$ (TVar Bool -> STM Bool
    forall a. TVar a -> STM a
    readTVar TVar Bool
    delayDoneVar STM Bool -> (Bool -> STM ()) -> STM ()
    forall a b. STM a -> (a -> STM b) -> STM b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= Bool -> STM ()
    checkSTM)) STM (mio ()) -> STM (mio ()) -> STM (mio ())
    forall a. STM a -> STM a -> STM a
    forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
    <|>
              -- Time to stop the companion, throw a 'CompanionDone' exception

              -- immediately

              (CompanionDone -> mio ()
    forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
    throwIO CompanionDone
    CompanionDone mio () -> STM () -> STM (mio ())
    forall a b. a -> STM b -> STM a
    forall (f :: * -> *) a b. Functor f => a -> f b -> f a
    <$ (TVar Bool -> STM Bool
    forall a. TVar a -> STM a
    readTVar TVar Bool
    shouldStopVar STM Bool -> (Bool -> STM ()) -> STM ()
    forall a b. STM a -> (a -> STM b) -> STM b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= Bool -> STM ()
    checkSTM))

      -- Run the 'Companion' and inner action together

      Concurrently m a -> m a
    forall (m :: * -> *) a. Concurrently m a -> m a
    runConcurrently (Concurrently m a -> m a) -> Concurrently m a -> m a
    forall a b. (a -> b) -> a -> b
    $
        -- Ignore a 'CompanionDone' exception from the companion, that's expected

        -- behavior

        StopCompanion m -> Concurrently m ()
    forall (m :: * -> *) a. m a -> Concurrently m a
    Concurrently (Companion m
    companion Int -> mio ()
    Delay
    delay StopCompanion m
    -> (CompanionDone -> StopCompanion m) -> StopCompanion m
    forall (m :: * -> *) e a.
    (MonadUnliftIO m, Exception e) =>
    m a -> (e -> m a) -> m a
    `catch` \CompanionDone
    CompanionDone -> () -> StopCompanion m
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()) Concurrently m () -> Concurrently m a -> Concurrently m a
    forall a b.
    Concurrently m a -> Concurrently m b -> Concurrently m b
    forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
    *>
        -- Run the inner action, giving it the 'StopCompanion' action, and

        -- ensuring it is called regardless of exceptions.

        m a -> Concurrently m a
    forall (m :: * -> *) a. m a -> Concurrently m a
    Concurrently (StopCompanion m -> m a
    inner StopCompanion m
    stopCompanion m a -> StopCompanion m -> m a
    forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
    `finally` StopCompanion m
    stopCompanion)