    {-# LANGUAGE CPP                 #-}
    {-# LANGUAGE TupleSections       #-}
    {-# LANGUAGE DeriveDataTypeable  #-}
    {-# LANGUAGE DeriveFunctor       #-}
    {-# LANGUAGE DeriveGeneric       #-}
    {-# LANGUAGE GADTs               #-}
    {-# LANGUAGE LambdaCase          #-}
    {-# LANGUAGE RankNTypes          #-}
    {-# LANGUAGE ScopedTypeVariables #-}
    {-# LANGUAGE StandaloneDeriving  #-}
    module UnliftIO.Internals.Async where

    import           Control.Applicative
    import           Control.Concurrent       (threadDelay, getNumCapabilities)
    import qualified Control.Concurrent       as C
    import           Control.Concurrent.Async (Async)
    import qualified Control.Concurrent.Async as A
    import           Control.Concurrent.STM
    import           Control.Exception        (Exception, SomeException)
    import           Control.Monad            (forever, liftM, unless, void, (>=>))
    import           Control.Monad.IO.Unlift
    import           Data.Foldable            (for_, traverse_)
    import           Data.Typeable            (Typeable)
    import           Data.IORef (IORef, readIORef, atomicWriteIORef, newIORef, atomicModifyIORef')
    import qualified UnliftIO.Exception       as UE

    -- For the implementation of Conc below, we do not want any of the
    -- smart async exception handling logic from UnliftIO.Exception, since
    -- (eg) we're low-level enough to need to explicit be throwing async
    -- exceptions synchronously.
    import qualified Control.Exception        as E
    import           GHC.Generics             (Generic)

    #if MIN_VERSION_base(4,9,0)
    import           Data.Semigroup
    #else
    import           Data.Monoid              hiding (Alt)
    #endif
    import           Data.Foldable            (Foldable, toList)
    import           Data.Traversable         (Traversable, for, traverse)

    -- | Unlifted 'A.async'.
    --
    -- @since 0.1.0.0
    async :: MonadUnliftIO m => m a -> m (Async a)
    async m = withRunInIO $ \run -> A.async $ run m

    -- | Unlifted 'A.asyncBound'.
    --
    -- @since 0.1.0.0
    asyncBound :: MonadUnliftIO m => m a -> m (Async a)
    asyncBound m = withRunInIO $ \run -> A.asyncBound $ run m

    -- | Unlifted 'A.asyncOn'.
    --
    -- @since 0.1.0.0
    asyncOn :: MonadUnliftIO m => Int -> m a -> m (Async a)
    asyncOn i m = withRunInIO $ \run -> A.asyncOn i $ run m

    -- | Unlifted 'A.asyncWithUnmask'.
    --
    -- @since 0.1.0.0
    asyncWithUnmask :: MonadUnliftIO m => ((forall b. m b -> m b) -> m a) -> m (Async a)
    asyncWithUnmask m =
      withRunInIO $ \run -> A.asyncWithUnmask $ \unmask -> run $ m $ liftIO . unmask . run

    -- | Unlifted 'A.asyncOnWithUnmask'.
    --
    -- @since 0.1.0.0
    asyncOnWithUnmask :: MonadUnliftIO m => Int -> ((forall b. m b -> m b) -> m a) -> m (Async a)
    asyncOnWithUnmask i m =
      withRunInIO $ \run -> A.asyncOnWithUnmask i $ \unmask -> run $ m $ liftIO . unmask . run

    -- | Unlifted 'A.withAsync'.
    --
    -- @since 0.1.0.0
    withAsync :: MonadUnliftIO m => m a -> (Async a -> m b) -> m b
    withAsync a b = withRunInIO $ \run -> A.withAsync (run a) (run . b)

    -- | Unlifted 'A.withAsyncBound'.
    --
    -- @since 0.1.0.0
    withAsyncBound :: MonadUnliftIO m => m a -> (Async a -> m b) -> m b
    withAsyncBound a b = withRunInIO $ \run -> A.withAsyncBound (run a) (run . b)

    -- | Unlifted 'A.withAsyncOn'.
    --
    -- @since 0.1.0.0
    withAsyncOn :: MonadUnliftIO m => Int -> m a -> (Async a -> m b) -> m b
    withAsyncOn i a b = withRunInIO $ \run -> A.withAsyncOn i (run a) (run . b)

    -- | Unlifted 'A.withAsyncWithUnmask'.
    --
    -- @since 0.1.0.0
    withAsyncWithUnmask
      :: MonadUnliftIO m
      => ((forall c. m c -> m c) -> m a)
      -> (Async a -> m b)
      -> m b
    withAsyncWithUnmask a b =
      withRunInIO $ \run -> A.withAsyncWithUnmask
        (\unmask -> run $ a $ liftIO . unmask . run)
        (run . b)

    -- | Unlifted 'A.withAsyncOnWithMask'.
    --
    -- @since 0.1.0.0
    withAsyncOnWithUnmask
      :: MonadUnliftIO m
      => Int
      -> ((forall c. m c -> m c) -> m a)
      -> (Async a -> m b)
      -> m b
    withAsyncOnWithUnmask i a b =
      withRunInIO $ \run -> A.withAsyncOnWithUnmask i
        (\unmask -> run $ a $ liftIO . unmask . run)
        (run . b)

    -- | Lifted 'A.wait'.
    --
    -- @since 0.1.0.0
    wait :: MonadIO m => Async a -> m a
    wait = liftIO . A.wait

    -- | Lifted 'A.poll'.
    --
    -- @since 0.1.0.0
    poll :: MonadIO m => Async a -> m (Maybe (Either SomeException a))
    poll = liftIO . A.poll

    -- | Lifted 'A.waitCatch'.
    --
    -- @since 0.1.0.0
    waitCatch :: MonadIO m => Async a -> m (Either SomeException a)
    waitCatch = liftIO . A.waitCatch

    -- | Lifted 'A.cancel'.
    --
    -- @since 0.1.0.0
    cancel :: MonadIO m => Async a -> m ()
    cancel = liftIO . A.cancel

    -- | Lifted 'A.uninterruptibleCancel'.
    --
    -- @since 0.1.0.0
    uninterruptibleCancel :: MonadIO m => Async a -> m ()
    uninterruptibleCancel = liftIO . A.uninterruptibleCancel

    -- | Lifted 'A.cancelWith'. Additionally uses 'UE.toAsyncException' to
    -- ensure async exception safety.
    --
    -- @since 0.1.0.0
    cancelWith :: (Exception e, MonadIO m) => Async a -> e -> m ()
    cancelWith a e = liftIO (A.cancelWith a (UE.toAsyncException e))

    -- | Lifted 'A.waitAny'.
    --
    -- @since 0.1.0.0
    waitAny :: MonadIO m => [Async a] -> m (Async a, a)
    waitAny = liftIO . A.waitAny

    -- | Lifted 'A.waitAnyCatch'.
    --
    -- @since 0.1.0.0
    waitAnyCatch :: MonadIO m => [Async a] -> m (Async a, Either SomeException a)
    waitAnyCatch = liftIO . A.waitAnyCatch

    -- | Lifted 'A.waitAnyCancel'.
    --
    -- @since 0.1.0.0
    waitAnyCancel :: MonadIO m => [Async a] -> m (Async a, a)
    waitAnyCancel = liftIO . A.waitAnyCancel

    -- | Lifted 'A.waitAnyCatchCancel'.
    --
    -- @since 0.1.0.0
    waitAnyCatchCancel :: MonadIO m => [Async a] -> m (Async a, Either SomeException a)
    waitAnyCatchCancel = liftIO . A.waitAnyCatchCancel

    -- | Lifted 'A.waitEither'.
    --
    -- @since 0.1.0.0
    waitEither :: MonadIO m => Async a -> Async b -> m (Either a b)
    waitEither a b = liftIO (A.waitEither a b)

    -- | Lifted 'A.waitEitherCatch'.
    --
    -- @since 0.1.0.0
    waitEitherCatch :: MonadIO m => Async a -> Async b -> m (Either (Either SomeException a) (Either SomeException b))
    waitEitherCatch a b = liftIO (A.waitEitherCatch a b)

    -- | Lifted 'A.waitEitherCancel'.
    --
    -- @since 0.1.0.0
    waitEitherCancel :: MonadIO m => Async a -> Async b -> m (Either a b)
    waitEitherCancel a b = liftIO (A.waitEitherCancel a b)

    -- | Lifted 'A.waitEitherCatchCancel'.
    --
    -- @since 0.1.0.0
    waitEitherCatchCancel :: MonadIO m => Async a -> Async b -> m (Either (Either SomeException a) (Either SomeException b))
    waitEitherCatchCancel a b = liftIO (A.waitEitherCatchCancel a b)

    -- | Lifted 'A.waitEither_'.
    --
    -- @since 0.1.0.0
    waitEither_ :: MonadIO m => Async a -> Async b -> m ()
    waitEither_ a b = liftIO (A.waitEither_ a b)

    -- | Lifted 'A.waitBoth'.
    --
    -- @since 0.1.0.0
    waitBoth :: MonadIO m => Async a -> Async b -> m (a, b)
    waitBoth a b = liftIO (A.waitBoth a b)

    -- | Lifted 'A.link'.
    --
    -- @since 0.1.0.0
    link :: MonadIO m => Async a -> m ()
    link = liftIO . A.link

    -- | Lifted 'A.link2'.
    --
    -- @since 0.1.0.0
    link2 :: MonadIO m => Async a -> Async b -> m ()
    link2 a b = liftIO (A.link2 a b)

    -- | Unlifted 'A.race'.
    --
    -- @since 0.1.0.0
    race :: MonadUnliftIO m => m a -> m b -> m (Either a b)
    race a b = withRunInIO $ \run -> A.race (run a) (run b)

    -- | Unlifted 'A.race_'.
    --
    -- @since 0.1.0.0
    race_ :: MonadUnliftIO m => m a -> m b -> m ()
    race_ a b = withRunInIO $ \run -> A.race_ (run a) (run b)

    -- | Unlifted 'A.concurrently'.
    --
    -- @since 0.1.0.0
    concurrently :: MonadUnliftIO m => m a -> m b -> m (a, b)
    concurrently a b = withRunInIO $ \run -> A.concurrently (run a) (run b)

    -- | Unlifted 'A.concurrently_'.
    --
    -- @since 0.1.0.0
    concurrently_ :: MonadUnliftIO m => m a -> m b -> m ()
    concurrently_ a b = withRunInIO $ \run -> A.concurrently_ (run a) (run b)

    -- | Unlifted 'A.Concurrently'.
    --
    -- @since 0.1.0.0
    newtype Concurrently m a = Concurrently
      { runConcurrently :: m a
      }

    -- | @since 0.1.0.0
    instance Monad m => Functor (Concurrently m) where
      fmap f (Concurrently a) = Concurrently $ liftM f a

    -- | @since 0.1.0.0
    instance MonadUnliftIO m => Applicative (Concurrently m) where
      pure = Concurrently . return
      Concurrently fs <*> Concurrently as =
        Concurrently $ liftM (\(f, a) -> f a) (concurrently fs as)

    -- | Composing two unlifted 'Concurrently' values using 'Alternative' is the
    -- equivalent to using a 'race' combinator, the asynchrounous sub-routine that
    -- returns a value first is the one that gets it's value returned, the slowest
    -- sub-routine gets cancelled and it's thread is killed.
    --
    -- @since 0.1.0.0
    instance MonadUnliftIO m => Alternative (Concurrently m) where
      -- | Care should be taken when using the 'empty' value of the 'Alternative'
      -- interface, as it will create a thread that delays for a long period of
      -- time. The reason behind this implementation is that any other computation
      -- will finish first than the 'empty' value. This implementation is less than
      -- ideal, and in a perfect world, we would have a typeclass family that allows
      -- '(<|>)' but not 'empty'.
      --
      -- @since 0.1.0.0
      empty = Concurrently $ liftIO (forever (threadDelay maxBound))
      Concurrently as <|> Concurrently bs =
        Concurrently $ liftM (either id id) (race as bs)

    --------------------------------------------------------------------------------
    #if MIN_VERSION_base(4,9,0)
    --------------------------------------------------------------------------------
    -- | Only defined by @async@ for @base >= 4.9@.
    --
    -- @since 0.1.0.0
    instance (MonadUnliftIO m, Semigroup a) => Semigroup (Concurrently m a) where
      (<>) = liftA2 (<>)

    -- | @since 0.1.0.0
    instance (Semigroup a, Monoid a, MonadUnliftIO m) => Monoid (Concurrently m a) where
      mempty = pure mempty
      mappend = (<>)
    --------------------------------------------------------------------------------
    #else
    --------------------------------------------------------------------------------
    -- | @since 0.1.0.0
    instance (Monoid a, MonadUnliftIO m) => Monoid (Concurrently m a) where
      mempty = pure mempty
      mappend = liftA2 mappend
    --------------------------------------------------------------------------------
    #endif
    --------------------------------------------------------------------------------

    -- | Similar to 'mapConcurrently' but with arguments flipped
    --
    -- @since 0.1.0.0
    forConcurrently :: MonadUnliftIO m => Traversable t => t a -> (a -> m b) -> m (t b)
    forConcurrently = flip mapConcurrently
    {-# INLINE forConcurrently #-}

    -- | Similar to 'mapConcurrently_' but with arguments flipped
    --
    -- @since 0.1.0.0
    forConcurrently_ :: MonadUnliftIO m => Foldable f => f a -> (a -> m b) -> m ()
    forConcurrently_ = flip mapConcurrently_
    {-# INLINE forConcurrently_ #-}

    -- | Unlifted 'A.replicateConcurrently'.
    --
    -- @since 0.1.0.0
    #if MIN_VERSION_base(4,7,0)
    #else
    replicateConcurrently :: (Functor m, MonadUnliftIO m) => Int -> m a -> m [a]
    #endif
    replicateConcurrently cnt m =
      case compare cnt 1 of
        LT -> pure []
        EQ -> (:[]) <$> m
        GT -> mapConcurrently id (replicate cnt m)
    {-# INLINE replicateConcurrently #-}

    -- | Unlifted 'A.replicateConcurrently_'.
    --
    -- @since 0.1.0.0
    #if MIN_VERSION_base(4,7,0)
    replicateConcurrently_ :: (Applicative m, MonadUnliftIO m) => Int -> m a -> m ()
    #else
    replicateConcurrently_ :: (MonadUnliftIO m) => Int -> m a -> m ()
    #endif
    replicateConcurrently_ cnt m =
      case compare cnt 1 of
        LT -> pure ()
        EQ -> void m
        GT -> mapConcurrently_ id (replicate cnt m)
    {-# INLINE replicateConcurrently_ #-}

    -- Conc uses GHC features that are not supported in versions <= to ghc-7.10
    -- so we are going to export/use it when we have a higher version only.
    --------------------------------------------------------------------------------
    #if MIN_VERSION_base(4,8,0)
    --------------------------------------------------------------------------------

    -- | Executes a 'Traversable' container of items concurrently, it uses the 'Flat'
    -- type internally.
    --
    -- @since 0.1.0.0
    mapConcurrently :: MonadUnliftIO m => Traversable t => (a -> m b) -> t a -> m (t b)
    mapConcurrently f t = withRunInIO $ \run -> runFlat $ traverse
      (FlatApp . FlatAction . run . f)
      t
    {-# INLINE mapConcurrently #-}

    -- | Executes a 'Traversable' container of items concurrently, it uses the 'Flat'
    -- type internally. This function ignores the results.
    --
    -- @since 0.1.0.0
    mapConcurrently_ :: MonadUnliftIO m => Foldable f => (a -> m b) -> f a -> m ()
    mapConcurrently_ f t = withRunInIO $ \run -> runFlat $ traverse_
      (FlatApp . FlatAction . run . f)
      t
    {-# INLINE mapConcurrently_ #-}

    -- More efficient Conc implementation

    -- | A more efficient alternative to 'Concurrently', which reduces the
    -- number of threads that need to be forked. For more information, see
    -- @FIXME link to blog post@. This is provided as a separate type to
    -- @Concurrently@ as it has a slightly different API.
    --
    -- Use the 'conc' function to construct values of type 'Conc', and
    -- 'runConc' to execute the composed actions. You can use the
    -- @Applicative@ instance to run different actions and wait for all of
    -- them to complete, or the @Alternative@ instance to wait for the
    -- first thread to complete.
    --
    -- In the event of a runtime exception thrown by any of the children
    -- threads, or an asynchronous exception received in the parent
    -- thread, all threads will be killed with an 'A.AsyncCancelled'
    -- exception and the original exception rethrown. If multiple
    -- exceptions are generated by different threads, there are no
    -- guarantees on which exception will end up getting rethrown.
    --
    -- For many common use cases, you may prefer using helper functions in
    -- this module like 'mapConcurrently'.
    --
    -- There are some intentional differences in behavior to
    -- @Concurrently@:
    --
    -- * Children threads are always launched in an unmasked state, not
    --   the inherited state of the parent thread.
    --
    -- Note that it is a programmer error to use the @Alternative@
    -- instance in such a way that there are no alternatives to an empty,
    -- e.g. @runConc (empty <|> empty)@. In such a case, a 'ConcException'
    -- will be thrown. If there was an @Alternative@ in the standard
    -- libraries without @empty@, this library would use it instead.
    --
    -- @since 0.2.9.0
    data Conc m a where
      Action :: m a -> Conc m a
      Apply   :: Conc m (v -> a) -> Conc m v -> Conc m a
      LiftA2 :: (x -> y -> a) -> Conc m x -> Conc m y -> Conc m a

      -- Just an optimization to avoid spawning extra threads
      Pure :: a -> Conc m a

      -- I thought there would be an optimization available from having a
      -- data constructor that explicit doesn't care about the first
      -- result. Turns out it doesn't help much: we still need to keep a
      -- TMVar below to know when the thread completes.
      --
      -- Then :: Conc m a -> Conc m b -> Conc m b

      Alt :: Conc m a -> Conc m a -> Conc m a
      Empty :: Conc m a

    deriving instance Functor m => Functor (Conc m)
    -- fmap f (Action routine) = Action (fmap f routine)
    -- fmap f (LiftA2 g x y)   = LiftA2 (fmap f g) x y
    -- fmap f (Pure val)       = Pure (f val)
    -- fmap f (Alt a b)        = Alt (fmap f a) (fmap f b)
    -- fmap f Empty            = Empty

    -- | Construct a value of type 'Conc' from an action. Compose these
    -- values using the typeclass instances (most commonly 'Applicative'
    -- and 'Alternative') and then run with 'runConc'.
    --
    -- @since 0.2.9.0
    conc :: m a -> Conc m a
    conc = Action
    {-# INLINE conc #-}

    -- | Run a 'Conc' value on multiple threads.
    --
    -- @since 0.2.9.0
    runConc :: MonadUnliftIO m => Conc m a -> m a
    runConc = flatten >=> (liftIO . runFlat)
    {-# INLINE runConc #-}

    -- | @since 0.2.9.0
    instance MonadUnliftIO m => Applicative (Conc m) where
      pure = Pure
      {-# INLINE pure #-}
      -- | Following is an example of how an 'Applicative' expands to a Tree
      --
      -- @@@
      -- downloadA :: IO String
      -- downloadB :: IO String
      --
      -- (f <$> conc downloadA <*> conc downloadB <*> pure 123)
      --
      --   (((f <$> a) <*> b) <*> c))
      --        (1)    (2)    (3)
      --
      -- (1)
      --   Action (fmap f downloadA)
      -- (2)
      --   Apply (Action (fmap f downloadA)) (Action downloadB)
      -- (3)
      --   Apply (Apply (Action (fmap f downloadA)) (Action downloadB))
      --        (Pure 123)
      -- @@@
      --
      (<*>) = Apply
      {-# INLINE (<*>) #-}
      -- See comment above on Then
      -- (*>) = Then
    #if MIN_VERSION_base(4,11,0)
      liftA2 = LiftA2
      {-# INLINE liftA2 #-}
    #endif

      a *> b = LiftA2 (\_ x -> x) a b
      {-# INLINE (*>) #-}

    -- | @since 0.2.9.0
    instance MonadUnliftIO m => Alternative (Conc m) where
      empty = Empty -- this is so ugly, we don't actually want to provide it!
      {-# INLINE empty #-}
      (<|>) = Alt
      {-# INLINE (<|>) #-}

    #if MIN_VERSION_base(4, 11, 0)
    -- | @since 0.2.9.0
    instance (MonadUnliftIO m, Semigroup a) => Semigroup (Conc m a) where
      (<>) = liftA2 (<>)
      {-# INLINE (<>) #-}
    #endif

    -- | @since 0.2.9.0
    instance (Monoid a, MonadUnliftIO m) => Monoid (Conc m a) where
      mempty = pure mempty
      {-# INLINE mempty #-}
      mappend = liftA2 mappend
      {-# INLINE mappend #-}

    -------------------------
    -- Conc implementation --
    -------------------------

    -- Data types for flattening out the original @Conc@ into a simplified
    -- view. Goals:
    --
    -- * We want to get rid of the Empty data constructor. We don't want
    --   it anyway, it's only there because of the Alternative typeclass.
    --
    -- * We want to ensure that there is no nesting of Alt data
    --   constructors. There is a bookkeeping overhead to each time we
    --   need to track raced threads, and we want to minimize that
    --   bookkeeping.
    --
    -- * We want to ensure that, when racing, we're always racing at least
    --   two threads.
    --
    -- * We want to simplify down to IO.

    -- | Flattened structure, either Applicative or Alternative
    data Flat a
      = FlatApp !(FlatApp a)
      -- | Flattened Alternative. Has at least 2 entries, which must be
      -- FlatApp (no nesting of FlatAlts).
      | FlatAlt !(FlatApp a) !(FlatApp a) ![FlatApp a]

    deriving instance Functor Flat
    -- fmap f (FlatApp a) =
    --  FlatApp (fmap f a)
    -- fmap f (FlatAlt (FlatApp a) (FlatApp b) xs) =
    --   FlatAlt (FlatApp (fmap f a)) (FlatApp (fmap f b)) (map (fmap f) xs)
    instance Applicative Flat where
      pure = FlatApp . pure
      (<*>) f a = FlatApp (FlatLiftA2 id f a)
    #if MIN_VERSION_base(4,11,0)
      liftA2 f a b = FlatApp (FlatLiftA2 f a b)
    #endif

    -- | Flattened Applicative. No Alternative stuff directly in here, but may be in
    -- the children. Notice this type doesn't have a type parameter for monadic
    -- contexts, it hardwires the base monad to IO given concurrency relies
    -- eventually on that.
    --
    -- @since 0.2.9.0
    data FlatApp a where
      FlatPure   :: a -> FlatApp a
      FlatAction :: IO a -> FlatApp a
      FlatApply   :: Flat (v -> a) -> Flat v -> FlatApp a
      FlatLiftA2 :: (x -> y -> a) -> Flat x -> Flat y -> FlatApp a

    deriving instance Functor FlatApp
    instance Applicative FlatApp where
      pure = FlatPure
      (<*>) mf ma = FlatApply (FlatApp mf) (FlatApp ma)
    #if MIN_VERSION_base(4,11,0)
      liftA2 f a b = FlatLiftA2 f (FlatApp a) (FlatApp b)
    #endif

    -- | Things that can go wrong in the structure of a 'Conc'. These are
    -- /programmer errors/.
    --
    -- @since 0.2.9.0
    data ConcException
      = EmptyWithNoAlternative
      deriving (Generic, Show, Typeable, Eq, Ord)
    instance E.Exception ConcException

    -- | Simple difference list, for nicer types below
    type DList a = [a] -> [a]

    dlistConcat :: DList a -> DList a -> DList a
    dlistConcat = (.)
    {-# INLINE dlistConcat #-}

    dlistCons :: a -> DList a -> DList a
    dlistCons a as = dlistSingleton a `dlistConcat` as
    {-# INLINE dlistCons #-}

    dlistConcatAll :: [DList a] -> DList a
    dlistConcatAll = foldr (.) id
    {-# INLINE dlistConcatAll #-}

    dlistToList :: DList a -> [a]
    dlistToList = ($ [])
    {-# INLINE dlistToList #-}

    dlistSingleton :: a -> DList a
    dlistSingleton a = (a:)
    {-# INLINE dlistSingleton #-}

    dlistEmpty :: DList a
    dlistEmpty = id
    {-# INLINE dlistEmpty #-}

    -- | Turn a 'Conc' into a 'Flat'. Note that thanks to the ugliness of
    -- 'empty', this may fail, e.g. @flatten Empty@.
    --
    -- @since 0.2.9.0
    flatten :: forall m a. MonadUnliftIO m => Conc m a -> m (Flat a)
    flatten c0 = withRunInIO $ \run -> do

      -- why not app?
      let both :: forall k. Conc m k -> IO (Flat k)
          both Empty = E.throwIO EmptyWithNoAlternative
          both (Action m) = pure $ FlatApp $ FlatAction $ run m
          both (Apply cf ca) = do
            f <- both cf
            a <- both ca
            pure $ FlatApp $ FlatApply f a
          both (LiftA2 f ca cb) = do
            a <- both ca
            b <- both cb
            pure $ FlatApp $ FlatLiftA2 f a b
          both (Alt ca cb) = do
            a <- alt ca
            b <- alt cb
            case dlistToList (a `dlistConcat` b) of
              []    -> E.throwIO EmptyWithNoAlternative
              [x]   -> pure $ FlatApp x
              x:y:z -> pure $ FlatAlt x y z
          both (Pure a) = pure $ FlatApp $ FlatPure a

          -- Returns a difference list for cheaper concatenation
          alt :: forall k. Conc m k -> IO (DList (FlatApp k))
          alt Empty = pure dlistEmpty
          alt (Apply cf ca) = do
            f <- both cf
            a <- both ca
            pure (dlistSingleton $ FlatApply f a)
          alt (Alt ca cb) = do
            a <- alt ca
            b <- alt cb
            pure $ a `dlistConcat` b
          alt (Action m) = pure (dlistSingleton $ FlatAction (run m))
          alt (LiftA2 f ca cb) = do
            a <- both ca
            b <- both cb
            pure (dlistSingleton $ FlatLiftA2 f a b)
          alt (Pure a) = pure (dlistSingleton $ FlatPure a)

      both c0

    -- | Run a @Flat a@ on multiple threads.
    runFlat :: Flat a -> IO a

    -- Silly, simple optimizations
    runFlat (FlatApp (FlatAction io)) = io
    runFlat (FlatApp (FlatPure x)) = pure x

    -- Start off with all exceptions masked so we can install proper cleanup.
    runFlat f0 = E.uninterruptibleMask $ \restore -> do
      -- How many threads have been spawned and finished their task? We need to
      -- ensure we kill all child threads and wait for them to die.
      resultCountVar <- newTVarIO 0

      -- Forks off as many threads as necessary to run the given Flat a,
      -- and returns:
      --
      -- + An STM action that will block until completion and return the
      --   result.
      --
      -- + The IDs of all forked threads. These need to be tracked so they
      --   can be killed (either when an exception is thrown, or when one
      --   of the alternatives completes first).
      --
      -- It would be nice to have the returned STM action return an Either
      -- and keep the SomeException values somewhat explicit, but in all
      -- my testing this absolutely kills performance. Instead, we're
      -- going to use a hack of providing a TMVar to fill up with a
      -- SomeException when things fail.
      --
      -- TODO: Investigate why performance degradation on Either
      let go :: forall a.
                TMVar E.SomeException
             -> Flat a
             -> IO (STM a, DList C.ThreadId)
          go _excVar (FlatApp (FlatPure x)) = pure (pure x, dlistEmpty)
          go excVar (FlatApp (FlatAction io)) = do
            resVar <- newEmptyTMVarIO
            tid <- C.forkIOWithUnmask $ \restore1 -> do
              res <- E.try $ restore1 io
              atomically $ do
                modifyTVar' resultCountVar (+ 1)
                case res of
                  Left e  -> void $ tryPutTMVar excVar e
                  Right x -> putTMVar resVar x
            pure (readTMVar resVar, dlistSingleton tid)
          go excVar (FlatApp (FlatApply cf ca)) = do
            (f, tidsf) <- go excVar cf
            (a, tidsa) <- go excVar ca
            pure (f <*> a, tidsf `dlistConcat` tidsa)
          go excVar (FlatApp (FlatLiftA2 f a b)) = do
            (a', tidsa) <- go excVar a
            (b', tidsb) <- go excVar b
            pure (liftA2 f a' b', tidsa `dlistConcat` tidsb)

          go excVar0 (FlatAlt x y z) = do
            -- As soon as one of the children finishes, we need to kill the siblings,
            -- we're going to create our own excVar here to pass to the children, so
            -- we can prevent the ThreadKilled exceptions we throw to the children
            -- here from propagating and taking down the whole system.
            excVar <- newEmptyTMVarIO
            resVar <- newEmptyTMVarIO
            pairs <- traverse (go excVar . FlatApp) (x:y:z)
            let (blockers, workerTids) = unzip pairs

            -- Fork a helper thread to wait for the first child to
            -- complete, or for one of them to die with an exception so we
            -- can propagate it to excVar0.
            helperTid <- C.forkIOWithUnmask $ \restore1 -> do
              eres <- E.try $ restore1 $ atomically $ foldr
                (\blocker rest -> (Right <$> blocker) <|> rest)
                (Left <$> readTMVar excVar)
                blockers
              atomically $ do
                modifyTVar' resultCountVar (+ 1)
                case eres of
                  -- NOTE: The child threads are spawned from @traverse go@ call above, they
                  -- are _not_ children of this helper thread, and helper thread doesn't throw
                  -- synchronous exceptions, so, any exception that the try above would catch
                  -- must be an async exception.
                  -- We were killed by an async exception, do nothing.
                  Left (_ :: E.SomeException) -> pure ()
                  -- Child thread died, propagate it
                  Right (Left e)              -> void $ tryPutTMVar excVar0 e
                  -- Successful result from one of the children
                  Right (Right res)           -> putTMVar resVar res

              -- And kill all of the threads
              for_ workerTids $ \tids' ->
                -- NOTE: Replacing A.AsyncCancelled with KillThread as the
                -- 'A.AsyncCancelled' constructor is not exported in older versions
                -- of the async package
                -- for_ (tids' []) $ \workerTid -> E.throwTo workerTid A.AsyncCancelled
                for_ (dlistToList tids') $ \workerTid -> C.killThread workerTid

            pure ( readTMVar resVar
                 , helperTid `dlistCons` dlistConcatAll workerTids
                 )

      excVar <- newEmptyTMVarIO
      (getRes, tids0) <- go excVar f0
      let tids = dlistToList tids0
          tidCount = length tids
          allDone count =
            if count > tidCount
              then error ("allDone: count ("
                          <> show count
                          <> ") should never be greater than tidCount ("
                          <> show tidCount
                          <> ")")
              else count == tidCount

      -- Automatically retry if we get killed by a
      -- BlockedIndefinitelyOnSTM. For more information, see:
      --
      -- + https:\/\/github.com\/simonmar\/async\/issues\/14
      -- + https:\/\/github.com\/simonmar\/async\/pull\/15
      --
      let autoRetry action =
            action `E.catch`
            \E.BlockedIndefinitelyOnSTM -> autoRetry action

      -- Restore the original masking state while blocking and catch
      -- exceptions to allow the parent thread to be killed early.
      res <- E.try $ restore $ autoRetry $ atomically $
             (Left <$> readTMVar excVar) <|>
             (Right <$> getRes)

      count0 <- atomically $ readTVar resultCountVar
      unless (allDone count0) $ do
        -- Kill all of the threads
        -- NOTE: Replacing A.AsyncCancelled with KillThread as the
        -- 'A.AsyncCancelled' constructor is not exported in older versions
        -- of the async package
        -- for_ tids $ \tid -> E.throwTo tid A.AsyncCancelled
        for_ tids $ \tid -> C.killThread tid

        -- Wait for all of the threads to die. We're going to restore the original
        -- masking state here, just in case there's a bug in the cleanup code of a
        -- child thread, so that we can be killed by an async exception. We decided
        -- this is a better behavior than hanging indefinitely and wait for a SIGKILL.
        restore $ atomically $ do
          count <- readTVar resultCountVar
          -- retries until resultCountVar has increased to the threadId count returned by go
          check $ allDone count

      -- Return the result or throw an exception. Yes, we could use
      -- either or join, but explicit pattern matching is nicer here.
      case res of
        -- Parent thread was killed with an async exception
        Left e          -> E.throwIO (e :: E.SomeException)
        -- Some child thread died
        Right (Left e)  -> E.throwIO e
        -- Everything worked!
        Right (Right x) -> pure x
    {-# INLINEABLE runFlat #-}

    --------------------------------------------------------------------------------
    #else
    --------------------------------------------------------------------------------

    -- | Unlifted 'A.mapConcurrently'.
    --
    -- @since 0.1.0.0
    mapConcurrently :: MonadUnliftIO m => Traversable t => (a -> m b) -> t a -> m (t b)
    mapConcurrently f t = withRunInIO $ \run -> A.mapConcurrently (run . f) t
    {-# INLINE mapConcurrently #-}

    -- | Unlifted 'A.mapConcurrently_'.
    --
    -- @since 0.1.0.0
    mapConcurrently_ :: MonadUnliftIO m => Foldable f => (a -> m b) -> f a -> m ()
    mapConcurrently_ f t = withRunInIO $ \run -> A.mapConcurrently_ (run . f) t
    {-# INLINE mapConcurrently_ #-}

    --------------------------------------------------------------------------------
    #endif
    --------------------------------------------------------------------------------

    -- | Like 'mapConcurrently' from async, but instead of one thread per
    -- element, it does pooling from a set of threads. This is useful in
    -- scenarios where resource consumption is bounded and for use cases
    -- where too many concurrent tasks aren't allowed.
    --
    -- === __Example usage__
    --
    -- @
    -- import Say
    --
    -- action :: Int -> IO Int
    -- action n = do
    --   tid <- myThreadId
    --   sayString $ show tid
    --   threadDelay (2 * 10^6) -- 2 seconds
    --   return n
    --
    -- main :: IO ()
    -- main = do
    --   yx \<- pooledMapConcurrentlyN 5 (\\x -\> action x) [1..5]
    --   print yx
    -- @
    --
    -- On executing you can see that five threads have been spawned:
    --
    -- @
    -- \$ ./pool
    -- ThreadId 36
    -- ThreadId 38
    -- ThreadId 40
    -- ThreadId 42
    -- ThreadId 44
    -- [1,2,3,4,5]
    -- @
    --
    --
    -- Let's modify the above program such that there are less threads
    -- than the number of items in the list:
    --
    -- @
    -- import Say
    --
    -- action :: Int -> IO Int
    -- action n = do
    --   tid <- myThreadId
    --   sayString $ show tid
    --   threadDelay (2 * 10^6) -- 2 seconds
    --   return n
    --
    -- main :: IO ()
    -- main = do
    --   yx \<- pooledMapConcurrentlyN 3 (\\x -\> action x) [1..5]
    --   print yx
    -- @
    -- On executing you can see that only three threads are active totally:
    --
    -- @
    -- \$ ./pool
    -- ThreadId 35
    -- ThreadId 37
    -- ThreadId 39
    -- ThreadId 35
    -- ThreadId 39
    -- [1,2,3,4,5]
    -- @
    --
    -- @since 0.2.10
    pooledMapConcurrentlyN :: (MonadUnliftIO m, Traversable t)
                          => Int -- ^ Max. number of threads. Should not be less than 1.
                          -> (a -> m b) -> t a -> m (t b)
    pooledMapConcurrentlyN numProcs f xs =
        withRunInIO $ \run -> pooledMapConcurrentlyIO numProcs (run . f) xs

    -- | Similar to 'pooledMapConcurrentlyN' but with number of threads
    -- set from 'getNumCapabilities'. Usually this is useful for CPU bound
    -- tasks.
    --
    -- @since 0.2.10
    pooledMapConcurrently :: (MonadUnliftIO m, Traversable t) => (a -> m b) -> t a -> m (t b)
    pooledMapConcurrently f xs = do
      withRunInIO $ \run -> do
        numProcs <- getNumCapabilities
        pooledMapConcurrentlyIO numProcs (run . f) xs

    -- | Similar to 'pooledMapConcurrentlyN' but with flipped arguments.
    --
    -- @since 0.2.10
    pooledForConcurrentlyN :: (MonadUnliftIO m, Traversable t)
                          => Int -- ^ Max. number of threads. Should not be less than 1.
                          -> t a -> (a -> m b) -> m (t b)
    pooledForConcurrentlyN numProcs = flip (pooledMapConcurrentlyN numProcs)

    -- | Similar to 'pooledForConcurrentlyN' but with number of threads
    -- set from 'getNumCapabilities'. Usually this is useful for CPU bound
    -- tasks.
    --
    -- @since 0.2.10
    pooledForConcurrently :: (MonadUnliftIO m, Traversable t) => t a -> (a -> m b) -> m (t b)
    pooledForConcurrently = flip pooledMapConcurrently

    pooledMapConcurrentlyIO :: Traversable t => Int -> (a -> IO b) -> t a -> IO (t b)
    pooledMapConcurrentlyIO numProcs f xs =
        if (numProcs < 1)
        then error "pooledMapconcurrentlyIO: number of threads < 1"
        else pooledMapConcurrentlyIO' numProcs f xs

    -- | Performs the actual pooling for the tasks. This function will
    -- continue execution until the task queue becomes empty. When one of
    -- the pooled thread finishes it's task, it will pickup the next task
    -- from the queue if an job is available.
    pooledConcurrently
      :: Int -- ^ Max. number of threads. Should not be less than 1.
      -> IORef [a] -- ^ Task queue. These are required as inputs for the jobs.
      -> (a -> IO ()) -- ^ The task which will be run concurrently (but
                     -- will be pooled properly).
      -> IO ()
    pooledConcurrently numProcs jobsVar f = do
      replicateConcurrently_ numProcs $ do
        let loop = do
              mbJob :: Maybe a <- atomicModifyIORef' jobsVar $ \x -> case x of
                [] -> ([], Nothing)
                var : vars -> (vars, Just var)
              case mbJob of
                Nothing -> return ()
                Just x -> do
                  f x
                  loop
         in loop

    pooledMapConcurrentlyIO' ::
        Traversable t => Int  -- ^ Max. number of threads. Should not be less than 1.
                      -> (a -> IO b)
                      -> t a
                      -> IO (t b)
    pooledMapConcurrentlyIO' numProcs f xs = do
      -- prepare one IORef per result...
      jobs :: t (a, IORef b) <-
        for xs (\x -> (x, ) <$> newIORef (error "pooledMapConcurrentlyIO': empty IORef"))
      -- ...put all the inputs in a queue..
      jobsVar :: IORef [(a, IORef b)] <- newIORef (toList jobs)
      -- ...run `numProcs` threads in parallel, each
      -- of them consuming the queue and filling in
      -- the respective IORefs.
      pooledConcurrently numProcs jobsVar $ \ (x, outRef) -> f x >>= atomicWriteIORef outRef      -- Read all the IORefs
      for jobs (\(_, outputRef) -> readIORef outputRef)

    pooledMapConcurrentlyIO_' ::
      Foldable t => Int -> (a -> IO ()) -> t a -> IO ()
    pooledMapConcurrentlyIO_' numProcs f jobs = do
      jobsVar :: IORef [a] <- newIORef (toList jobs)
      pooledConcurrently numProcs jobsVar f

    pooledMapConcurrentlyIO_ :: Foldable t => Int -> (a -> IO b) -> t a -> IO ()
    pooledMapConcurrentlyIO_ numProcs f xs =
        if (numProcs < 1)
        then error "pooledMapconcurrentlyIO_: number of threads < 1"
        else pooledMapConcurrentlyIO_' numProcs (\x -> f x >> return ()) xs

    -- | Like 'pooledMapConcurrentlyN' but with the return value
    -- discarded.
    --
    -- @since 0.2.10
    pooledMapConcurrentlyN_ :: (MonadUnliftIO m, Foldable f)
                            => Int -- ^ Max. number of threads. Should not be less than 1.
                            -> (a -> m b) -> f a -> m ()
    pooledMapConcurrentlyN_ numProcs f t =
      withRunInIO $ \run -> pooledMapConcurrentlyIO_ numProcs (run . f) t

    -- | Like 'pooledMapConcurrently' but with the return value discarded.
    --
    -- @since 0.2.10
    pooledMapConcurrently_ :: (MonadUnliftIO m, Foldable f) => (a -> m b) -> f a -> m ()
    pooledMapConcurrently_ f t =
      withRunInIO $ \run -> do
        numProcs <- getNumCapabilities
        pooledMapConcurrentlyIO_ numProcs (run . f) t

    -- | Like 'pooledMapConcurrently_' but with flipped arguments.
    --
    -- @since 0.2.10
    pooledForConcurrently_ :: (MonadUnliftIO m, Foldable f) => f a -> (a -> m b) -> m ()
    pooledForConcurrently_ = flip pooledMapConcurrently_

    -- | Like 'pooledMapConcurrentlyN_' but with flipped arguments.
    --
    -- @since 0.2.10
    pooledForConcurrentlyN_ :: (MonadUnliftIO m, Foldable t)
                            => Int -- ^ Max. number of threads. Should not be less than 1.
                            -> t a -> (a -> m b) -> m ()
    pooledForConcurrentlyN_ numProcs = flip (pooledMapConcurrentlyN_ numProcs)

    -- | Pooled version of 'replicateConcurrently'. Performs the action in
    -- the pooled threads.
    --
    -- @since 0.2.10
    pooledReplicateConcurrentlyN :: (MonadUnliftIO m)
                                 => Int -- ^ Max. number of threads. Should not be less than 1.
                                 -> Int -- ^ Number of times to perform the action.
                                 -> m a -> m [a]
    pooledReplicateConcurrentlyN numProcs cnt task =
        if cnt < 1
        then return []
        else pooledMapConcurrentlyN numProcs (\_ -> task) [1..cnt]

    -- | Similar to 'pooledReplicateConcurrentlyN' but with number of
    -- threads set from 'getNumCapabilities'. Usually this is useful for
    -- CPU bound tasks.
    --
    -- @since 0.2.10
    pooledReplicateConcurrently :: (MonadUnliftIO m)
                                => Int -- ^ Number of times to perform the action.
                                -> m a -> m [a]
    pooledReplicateConcurrently cnt task =
        if cnt < 1
        then return []
        else pooledMapConcurrently (\_ -> task) [1..cnt]

    -- | Pooled version of 'replicateConcurrently_'. Performs the action in
    -- the pooled threads.
    --
    -- @since 0.2.10
    pooledReplicateConcurrentlyN_ :: (MonadUnliftIO m)
                                  => Int -- ^ Max. number of threads. Should not be less than 1.
                                  -> Int -- ^ Number of times to perform the action.
                                  -> m a -> m ()
    pooledReplicateConcurrentlyN_ numProcs cnt task =
      if cnt < 1
      then return ()
      else pooledMapConcurrentlyN_ numProcs (\_ -> task) [1..cnt]

    -- | Similar to 'pooledReplicateConcurrently_' but with number of
    -- threads set from 'getNumCapabilities'. Usually this is useful for
    -- CPU bound tasks.
    --
    -- @since 0.2.10
    pooledReplicateConcurrently_ :: (MonadUnliftIO m)
                                 => Int -- ^ Number of times to perform the action.
                                 -> m a -> m ()
    pooledReplicateConcurrently_ cnt task =
      if cnt < 1
      then return ()
      else pooledMapConcurrently_ (\_ -> task) [1..cnt]