module UnliftIO.Async.Utils
  ( allocateAsync, allocateLinkedAsync
  , allocateAsyncWithUnmask, allocateLinkedAsyncWithUnmask
  , allocateAsyncMasked, allocateLinkedAsyncMasked
  ) where

import ClassyPrelude hiding (cancel, async, link)
import Control.Lens

import qualified UnliftIO.Async as UnliftIO
import qualified Control.Concurrent.Async as A

import Control.Monad.Trans.Resource


allocateAsync :: forall m a.
                 ( MonadUnliftIO m, MonadResource m )
              => m a -> m (Async a)
allocateAsync act = withRunInIO $ \run -> run . fmap (view _2) . flip allocate A.cancel . A.async $ run act

allocateLinkedAsync :: forall m a. (MonadUnliftIO m, MonadResource m) => m a -> m (Async a)
allocateLinkedAsync = uncurry (<$) . (id &&& UnliftIO.link) <=< allocateAsync


allocateAsyncWithUnmask :: forall m a.
                           ( MonadUnliftIO m, MonadResource m )
                        => ((forall b. m b -> m b) -> m a) -> m (Async a)
allocateAsyncWithUnmask act = withRunInIO $ \run -> run . fmap (view _2) . flip allocate A.cancel $ A.asyncWithUnmask $ \unmask -> run $ act (liftIO . unmask . run)

allocateLinkedAsyncWithUnmask :: forall m a. (MonadUnliftIO m, MonadResource m) => ((forall b. m b -> m b) -> m a) -> m (Async a)
allocateLinkedAsyncWithUnmask act = uncurry (<$) . (id &&& UnliftIO.link) =<< allocateAsyncWithUnmask act


allocateAsyncMasked :: forall m a.
                       ( MonadUnliftIO m, MonadResource m )
                    => m a -> m (Async a)
allocateAsyncMasked act = allocateAsyncWithUnmask (const act)

allocateLinkedAsyncMasked :: forall m a. (MonadUnliftIO m, MonadResource m) => m a -> m (Async a)
allocateLinkedAsyncMasked act = allocateLinkedAsyncWithUnmask (const act)