content: module Ki.Internal.Scope
    ( Scope,
      scoped,
      awaitAll,
      fork,
      forkWith,
      forkWith_,
      fork_,
      forkTry,
      forkTryWith,
    )
  where
  
  import Control.Concurrent (ThreadId, myThreadId, throwTo)
  import Control.Concurrent.MVar (MVar, newEmptyMVar, tryPutMVar, tryTakeMVar)
  import Control.Exception
    ( Exception (fromException, toException),
      MaskingState (..),
      SomeAsyncException,
      SomeException,
      asyncExceptionFromException,
      asyncExceptionToException,
      throwIO,
      try,
      uninterruptibleMask,
      pattern ErrorCall,
    )
  import Control.Monad (guard, when)
  import Data.Foldable (for_)
  import Data.Functor (void)
  import Data.IntMap (IntMap)
  import qualified Data.IntMap.Lazy as IntMap.Lazy
  import Data.Void (Void, absurd)
  import GHC.Conc
    ( STM,
      TVar,
      atomically,
      enableAllocationLimit,
      labelThread,
      newTVarIO,
      readTVar,
      retry,
      setAllocationCounter,
      throwSTM,
      writeTVar,
    )
  import GHC.Conc.Sync (readTVarIO)
  import GHC.IO (unsafeUnmask)
  import IntSupply (IntSupply)
  import qualified IntSupply
  import Ki.Internal.ByteCount (byteCountToInt64)
  import Ki.Internal.IO
    ( IOResult (..),
      UnexceptionalIO (..),
      assertM,
      exceptionIs,
      interruptiblyMasked,
      unexceptionalTry,
      unexceptionalTryEither,
      uninterruptiblyMasked,
    )
  import Ki.Internal.NonblockingSTM
  import Ki.Internal.Propagating (Tid, peelOffPropagating, propagate)
  import Ki.Internal.Thread (Thread, makeThread)
  import Ki.Internal.ThreadAffinity (forkWithAffinity)
  import Ki.Internal.ThreadOptions (ThreadOptions (..), defaultThreadOptions)
  
  -- | A scope.
  --
  -- ==== __👉 Details__
  --
  -- * A scope delimits the lifetime of all threads created within it.
  --
  -- * A scope is only valid during the callback provided to 'Ki.scoped'.
  --
  -- * The thread that creates a scope is considered the parent of all threads created within it.
  --
  -- * All threads created within a scope can be awaited together (see 'Ki.awaitAll').
  --
  -- * All threads created within a scope are terminated when the scope closes.
  data Scope = Scope
    { -- The MVar that a child tries to put to, in the case that it tries to propagate an exception to its parent, but
      -- gets delivered an exception from its parent concurrently (which interrupts the throw). The parent must raise
      -- exceptions in its children with asynchronous exceptions uninterruptibly masked for correctness, yet we don't want
      -- a parent in the process of tearing down to miss/ignore this exception that we're trying to propagate?
      --
      -- Why a single-celled MVar? What if two siblings are fighting to inform their parent of their death? Well, only
      -- one exception can be propagated by the parent anyway, so we wouldn't need or want both.
      childExceptionVar :: {-# UNPACK #-} !(MVar SomeException),
      -- The set of child threads that are currently running, each keyed by a monotonically increasing int.
      childrenVar :: {-# UNPACK #-} !(TVar (IntMap ThreadId)),
      -- The supply that holds the (int) key to use for the next child thread.
      nextChildIdSupply :: {-# UNPACK #-} !IntSupply,
      -- The id of the thread that created the scope, which is considered the parent of all threads created within it.
      parentThreadId :: {-# UNPACK #-} !ThreadId,
      statusVar :: {-# UNPACK #-} !(TVar ScopeStatus)
    }
  
  -- The scope status: either open (allowing new threads to be created), closing (disallowing new threads to be
  -- created, and in the process of killing running children), or closed (at the very end of `scoped`)
  type ScopeStatus = Int
  
  -- The number of child threads that are guaranteed to be about to start, in the sense that only the GHC scheduler
  -- can continue to delay; there's no opportunity for an async exception to strike and prevent one of these threads
  -- from starting.
  pattern Open :: Int
  pattern Open <- ((>= 0) -> True)
  
  -- The scope is closing.
  pattern Closing :: Int
  pattern Closing = -1
  
  -- The scope is closed.
  pattern Closed :: Int
  pattern Closed = -2
  
  {-# COMPLETE Open, Closing, Closed #-}
  
  -- Internal async exception thrown by a parent thread to its children when the scope is closing.
  --
  -- In various places we trust without verifying that any 'ScopeClosing' exception, which is not exported by this module,
  -- was indeed thrown to a thread by its parent. It is possible to write a program that violates this (just catch the
  -- async exception and throw it to some other thread)... but who would do that?
  data ScopeClosing
    = ScopeClosing
  
  instance Show ScopeClosing where
    show _ = "<<internal ki exception: scope closing>>"
  
  instance Exception ScopeClosing where
    toException = asyncExceptionToException
    fromException = asyncExceptionFromException
  
  -- | Open a scope, perform an IO action with it, then close the scope.
  --
  -- ==== __👉 Details__
  --
  -- * The thread that creates a scope is considered the parent of all threads created within it.
  --
  -- * A scope is only valid during the callback provided to 'Ki.scoped'.
  --
  -- * When a scope closes (/i.e./ just before 'Ki.scoped' returns):
  --
  --     * The parent thread raises an exception in all of its living children.
  --     * The parent thread blocks until those threads terminate.
  scoped :: (Scope -> IO a) -> IO a
  scoped action = do
    scope@Scope {childExceptionVar, childrenVar, statusVar} <- allocateScope
  
    uninterruptibleMask \restore -> do
      result <- try (restore (action scope))
  
      !runningChildren <- do
        atomically do
          -- Block until we haven't committed to starting any threads. Without this, we may create a thread concurrently
          -- with closing its scope, and not grab its thread id to throw an exception to.
          starting <- readTVar statusVar
          assertM (starting >= 0)
          guard (starting == 0)
          -- Indicate that this scope is closing, so attempts to create a new thread within it will throw ScopeClosing
          -- (as if the calling thread was a parent of this scope, which it should be, and we threw it a ScopeClosing
          -- ourselves).
          writeTVar statusVar Closing
          -- Return the list of currently-running children to kill. Some of them may have *just* started (e.g. if we
          -- initially retried in `guard (n == 0)` above). That's fine - kill them all!
          readTVar childrenVar
  
      -- Deliver a ScopeClosing exception to every running child.
      --
      -- This happens to throw in the order the children were created, but that isn't an important/useful enough feature
      -- to be worth documenting, so users shouldn't rely on it. It's definitely not the case that child 1 will completely
      -- terminate before child 2 is delivered an exception: each child may delay arbitrarily while cleaning up.
      for_ runningChildren \child -> throwTo child ScopeClosing
  
      atomically do
        -- Block until all children have terminated; this relies on children respecting the async exception, which they
        -- must, for correctness. Otherwise, a thread could indeed outlive the scope in which it's created, which is
        -- definitely not structured concurrency!
        children <- readTVar childrenVar
        guard (IntMap.Lazy.null children)
  
        -- Record the scope as closed (from closing), so subsequent attempts to use it will throw a runtime exception
        writeTVar statusVar Closed
  
      -- By now there are three sources of exception:
      --
      --   1) A sync or async exception thrown during the callback, captured in `result`. If applicable, we want to peel
      --      the `Propagating` off of this, which was only used to indicate it came from one of our children.
      --
      --   2) A sync or async exception left for us in `childExceptionVar` by a child that tried to propagate it to us
      --      directly, but failed (because we killed it concurrently).
      --
      --   3) An async exception waiting in our exception queue, because we still have async exceptions uninterruptibly
      --      masked.
      --
      -- We cannot throw more than one, so throw them in that priority order.
      case result of
        Left exception -> throwIO (peelOffPropagating exception)
        Right value ->
          tryTakeMVar childExceptionVar >>= \case
            Nothing -> pure value
            Just exception -> throwIO exception
  
  -- Allocate a new scope.
  allocateScope :: IO Scope
  allocateScope = do
    childExceptionVar <- newEmptyMVar
    childrenVar <- newTVarIO IntMap.Lazy.empty
    nextChildIdSupply <- IntSupply.new
    parentThreadId <- myThreadId
    statusVar <- newTVarIO 0
    pure Scope {childExceptionVar, childrenVar, nextChildIdSupply, parentThreadId, statusVar}
  
  -- Spawn a thread in a scope, providing it its child id and a function that sets the masking state to the requested
  -- masking state. The given action is called with async exceptions interruptibly masked.
  spawn :: Scope -> ThreadOptions -> (Tid -> (forall x. IO x -> IO x) -> UnexceptionalIO ()) -> IO ChildIds
  spawn scope@Scope {childrenVar, statusVar} options action = do
    -- Interruptible mask is enough so long as none of the STM operations below block.
    --
    -- Unconditionally set masking state to MaskedInterruptible, even though we might already be at MaskedInterruptible
    -- or MaskedUninterruptible, to avoid a branch on parentMaskingState.
    interruptiblyMasked do
      -- Record the thread as being about to start. Not allowed to retry.
      nonblockingAtomically do
        status <- nonblockingReadTVar statusVar
        assertM (status >= -2)
        case status of
          Open -> nonblockingWriteTVar' statusVar (status + 1)
          Closing -> nonblockingThrowSTM ScopeClosing
          Closed -> nonblockingThrowSTM (ErrorCall "ki: scope closed")
  
      childIds <- spawnChild scope options action
  
      -- Record the child as having started. Not allowed to retry.
      nonblockingAtomically do
        starting <- nonblockingReadTVar statusVar
        assertM (starting >= 1)
        nonblockingWriteTVar' statusVar (starting - 1)
        recordChild childrenVar childIds
  
      pure childIds
  
  data ChildIds
    = ChildIds
        {-# UNPACK #-} !Tid
        {-# UNPACK #-} !ThreadId
  
  spawnChild :: Scope -> ThreadOptions -> (Tid -> (forall x. IO x -> IO x) -> UnexceptionalIO ()) -> IO ChildIds
  spawnChild scope options action = do
    childId <- IntSupply.next nextChildIdSupply
    childThreadId <-
      forkWithAffinity affinity do
        when (not (null label)) do
          childThreadId <- myThreadId
          labelThread childThreadId label
  
        for_ allocationLimit \bytes -> do
          setAllocationCounter (byteCountToInt64 bytes)
          enableAllocationLimit
  
        let -- Action that sets the masking state from the current (MaskedInterruptible) to the requested one.
            atRequestedMaskingState :: IO a -> IO a
            atRequestedMaskingState =
              case requestedChildMaskingState of
                Unmasked -> unsafeUnmask
                MaskedInterruptible -> id
                MaskedUninterruptible -> uninterruptiblyMasked
  
        runUnexceptionalIO (action childId atRequestedMaskingState)
  
        nonblockingAtomically (unrecordChild childrenVar childId)
    pure (ChildIds childId childThreadId)
    where
      Scope {childrenVar, nextChildIdSupply} = scope
      ThreadOptions {affinity, allocationLimit, label, maskingState = requestedChildMaskingState} = options
  {-# INLINE spawnChild #-}
  
  -- Record our child by either:
  --
  --   * Flipping `Nothing` to `Just childThreadId` (common case: we record child before it unrecords itself)
  --   * Flipping `Just _` to `Nothing` (uncommon case: we observe that a child already unrecorded itself)
  recordChild :: TVar (IntMap ThreadId) -> ChildIds -> NonblockingSTM ()
  recordChild childrenVar (ChildIds childId childThreadId) = do
    children <- nonblockingReadTVar childrenVar
    nonblockingWriteTVar' childrenVar (IntMap.Lazy.alter (maybe (Just childThreadId) (const Nothing)) childId children)
  
  -- Unrecord a child (ourselves) by either:
  --
  --   * Flipping `Just childThreadId` to `Nothing` (common case: parent recorded us first)
  --   * Flipping `Nothing` to `Just undefined` (uncommon case: we terminate and unrecord before parent can record us).
  unrecordChild :: TVar (IntMap ThreadId) -> Tid -> NonblockingSTM ()
  unrecordChild childrenVar childId = do
    children <- nonblockingReadTVar childrenVar
    nonblockingWriteTVar' childrenVar (IntMap.Lazy.alter (maybe (Just undefined) (const Nothing)) childId children)
  
  -- | Wait until all threads created within a scope terminate.
  awaitAll :: Scope -> STM ()
  awaitAll Scope {childrenVar, statusVar} = do
    children <- readTVar childrenVar
    guard (IntMap.Lazy.null children)
    status <- readTVar statusVar
    assertM (status >= -2)
    case status of
      Open -> guard (status == 0)
      Closing -> retry -- block until closed
      Closed -> pure ()
  
  -- | Create a child thread to execute an action within a scope.
  --
  -- /Note/: The child thread does not mask asynchronous exceptions, regardless of the parent thread's masking state. To
  -- create a child thread with a different initial masking state, use 'Ki.forkWith'.
  fork :: Scope -> IO a -> IO (Thread a)
  fork scope =
    forkWith scope defaultThreadOptions
  
  -- | Variant of 'Ki.fork' for threads that never return.
  fork_ :: Scope -> IO Void -> IO ()
  fork_ scope =
    forkWith_ scope defaultThreadOptions
  
  -- | Variant of 'Ki.fork' that takes an additional options argument.
  forkWith :: Scope -> ThreadOptions -> IO a -> IO (Thread a)
  forkWith scope opts action = do
    resultVar <- newTVarIO NoResultYet
    let done result = UnexceptionalIO (atomically (writeTVar resultVar result))
    ChildIds _ childThreadId <-
      spawn scope opts \childId masking -> do
        unexceptionalTry (masking action) >>= \case
          Failure exception -> do
            when (not (exceptionIs @ScopeClosing exception)) do
              propagateException scope childId exception
            -- even put async exceptions that we propagated. this isn't totally ideal because a caller awaiting this
            -- thread would not be able to distinguish between async exceptions delivered to this thread, or itself
            done (BadResult exception)
          Success value -> done (GoodResult value)
    let doAwait =
          readTVar resultVar >>= \case
            NoResultYet -> retry
            BadResult exception -> throwSTM exception
            GoodResult value -> pure value
    pure (makeThread childThreadId doAwait)
  
  -- | Variant of 'Ki.forkWith' for threads that never return.
  forkWith_ :: Scope -> ThreadOptions -> IO Void -> IO ()
  forkWith_ scope opts action = do
    _childThreadId <-
      spawn scope opts \childId masking ->
        unexceptionalTryEither
          ( \exception ->
              when (not (exceptionIs @ScopeClosing exception)) do
                propagateException scope childId exception
          )
          absurd
          (masking action)
    pure ()
  
  -- | Like 'Ki.fork', but the child thread does not propagate exceptions that are both:
  --
  -- * Synchronous (/i.e./ not an instance of 'SomeAsyncException').
  -- * An instance of @e@.
  forkTry :: forall e a. (Exception e) => Scope -> IO a -> IO (Thread (Either e a))
  forkTry scope =
    forkTryWith scope defaultThreadOptions
  
  data Result a
    = NoResultYet
    | BadResult !SomeException -- sync or async
    | GoodResult a
  
  -- | Variant of 'Ki.forkTry' that takes an additional options argument.
  forkTryWith :: forall e a. (Exception e) => Scope -> ThreadOptions -> IO a -> IO (Thread (Either e a))
  forkTryWith scope opts action = do
    resultVar <- newTVarIO NoResultYet
    let done result = UnexceptionalIO (atomically (writeTVar resultVar result))
    ChildIds _ childThreadId <-
      spawn scope opts \childId masking -> do
        result <- unexceptionalTry (masking action)
        case result of
          Failure exception -> do
            -- then-branch explanation: if the user calls `forkTry @MyAsyncException` for some reason, we want to ignore
            -- this request and propagate the async exception. `forkTry` can only be used to catch synchronous exceptions.
            let shouldPropagate =
                  if exceptionIs @e exception
                    then exceptionIs @SomeAsyncException exception
                    else not (exceptionIs @ScopeClosing exception)
            when shouldPropagate (propagateException scope childId exception)
            done (BadResult exception)
          Success value -> done (GoodResult value)
    let doAwait =
          readTVar resultVar >>= \case
            NoResultYet -> retry
            BadResult exception ->
              case fromException @e exception of
                Nothing -> throwSTM exception
                Just expectedException -> pure (Left expectedException)
            GoodResult value -> pure (Right value)
    pure (makeThread childThreadId doAwait)
  
  -- We have a non-`ScopeClosing` exception to propagate to our parent.
  --
  -- If our scope has already begun closing (`statusVar` is Closing), then either...
  --
  --   (A) We already received a `ScopeClosing`, but then ended up trying to propagate an exception anyway, because we
  --   threw a synchronous exception (or were hit by a different asynchronous exception) during our teardown procedure.
  --
  --   or
  --
  --   (B) We will receive a `ScopeClosing` imminently, because our parent has *just* finished setting `statusVar` to
  --   Closing, and will proceed to throw ScopeClosing to all of its children.
  --
  -- If (A), our parent has asynchronous exceptions masked, so we must inform it of our exception via `childExceptionVar`
  -- rather than throwTo. If (B), either mechanism would work. And because we don't if we're in case (A) or (B), we just
  -- `childExceptionVar`.
  --
  -- And if our scope has not already begun closing (`statusVar` is not Closing), then we ought to throw our exception to
  -- it. But that might fail due to either...
  --
  --   (C) Our parent concurrently closing the scope and sending us a `ScopeClosing`; because it has asynchronous
  --   exceptions uninterruptibly masked and we only have asynchronous exception *synchronously* masked, its `throwTo`
  --   will return `()`, and ours will throw that `ScopeClosing` asynchronous exception. In this case, since we now know
  --   our parent is tearing down and has asynchronous exceptions masked, we again inform it via `childExceptionVar`.
  --
  --   (D) Some *other* non-`ScopeClosing` asynchronous exception is raised here. This is truly odd: maybe it's a heap
  --   overflow exception from the GHC runtime? Maybe some other thread has smuggled our `ThreadId` out and has manually
  --   thrown us an exception for some reason? Either way, because we already have an exception that we are trying to
  --   propagate, we just scoot these freaky exceptions under the rug.
  --
  -- Precondition: interruptibly masked
  propagateException :: Scope -> Tid -> SomeException -> UnexceptionalIO ()
  propagateException Scope {childExceptionVar, parentThreadId, statusVar} childId exception =
    UnexceptionalIO (readTVarIO statusVar) >>= \case
      Closing -> tryPutChildExceptionVar -- (A) or (B), we don't care which
      status -> do
        assertM (status >= 0) -- we know status is Open (0+) here; can't be Closed (-2)
        loop
    where
      loop :: UnexceptionalIO ()
      loop =
        unexceptionalTry (propagate exception childId parentThreadId) >>= \case
          Failure secondException
            | exceptionIs @ScopeClosing secondException -> tryPutChildExceptionVar -- (C)
            | otherwise -> loop -- (D)
          Success _ -> pure ()
  
      tryPutChildExceptionVar :: UnexceptionalIO ()
      tryPutChildExceptionVar =
        UnexceptionalIO (void (tryPutMVar childExceptionVar exception))
  