IO.hs

{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}

-- | Miscellaneous IO utilities
module Ki.Internal.IO
  ( -- * Unexceptional IO
    UnexceptionalIO (..),
    IOResult (..),
    unexceptionalTry,
    unexceptionalTryEither,

    -- * Exception utils
    assertIO,
    assertM,
    exceptionIs,
    interruptiblyMasked,
    uninterruptiblyMasked,
    tryEitherSTM,

    -- * Fork utils
    forkIO,
    forkOn,
  )
where

import Control.Exception
import Control.Monad (join)
import Data.Coerce (coerce)
import Data.Maybe (isJust)
import GHC.Base (maskAsyncExceptions#, maskUninterruptible#)
import GHC.Conc (STM, ThreadId (ThreadId), catchSTM)
import GHC.Exts (Int (I#), fork#, forkOn#)
import GHC.IO (IO (IO))
import System.IO.Unsafe (unsafePerformIO)
import Prelude

-- A little promise that this IO action cannot throw an exception (*including* async exceptions, which you normally
-- think of as being able to strike at any time).
--
-- Yeah it's verbose, and maybe not that necessary, but the code that bothers to use it really does require
-- un-exceptiony IO actions for correctness, so here we are.
newtype UnexceptionalIO a = UnexceptionalIO
  {runUnexceptionalIO :: IO a}
  deriving newtype (Applicative, Functor, Monad)

data IOResult a
  = Failure !SomeException -- sync or async exception
  | Success a

-- Try an action, catching any exception it throws.
--
-- The caller is responsible for ensuring that async exceptions are masked (at whatever masking level is appropriate),
-- as (again) `UnexceptionalIO` implies async exceptions won't be thrown either.
unexceptionalTry :: forall a. IO a -> UnexceptionalIO (IOResult a)
unexceptionalTry action =
  UnexceptionalIO do
    (Success <$> action) `catch` \exception ->
      pure (Failure exception)

-- Like try, but with continuations.
unexceptionalTryEither ::
  forall a b.
  (SomeException -> UnexceptionalIO b) ->
  (a -> UnexceptionalIO b) ->
  IO a ->
  UnexceptionalIO b
unexceptionalTryEither onFailure onSuccess action =
  UnexceptionalIO do
    join do
      catch
        (coerce @_ @(a -> IO b) onSuccess <$> action)
        (pure . coerce @_ @(SomeException -> IO b) onFailure)

-- | Make an assertion in a IO that requires IO.
assertIO :: IO Bool -> IO ()
assertIO b =
  assert (unsafePerformIO b) (pure ())
{-# INLINE assertIO #-}

-- | Make an assertion in a monad.
assertM :: (Applicative m) => Bool -> m ()
assertM b =
  assert b (pure ())
{-# INLINE assertM #-}

-- | @exceptionIs \@e exception@ returns whether @exception@ is an instance of @e@.
exceptionIs :: forall e. (Exception e) => SomeException -> Bool
exceptionIs =
  isJust . fromException @e

-- | Call an action with asynchronous exceptions interruptibly masked.
interruptiblyMasked :: forall a. IO a -> IO a
interruptiblyMasked =
  coerce (maskAsyncExceptions# @a)

-- | Call an action with asynchronous exceptions uninterruptibly masked.
uninterruptiblyMasked :: forall a. IO a -> IO a
uninterruptiblyMasked =
  coerce (maskUninterruptible# @a)

-- Like try, but with continuations
tryEitherSTM :: (Exception e) => (e -> STM b) -> (a -> STM b) -> STM a -> STM b
tryEitherSTM onFailure onSuccess action =
  join (catchSTM (onSuccess <$> action) (pure . onFailure))

-- Control.Concurrent.forkIO without the exception handler
forkIO :: IO () -> IO ThreadId
forkIO (IO action) =
  IO \s0 ->
    case fork# action s0 of
      (# s1, tid #) -> (# s1, ThreadId tid #)

-- Control.Concurrent.forkOn without the exception handler
forkOn :: Int -> IO () -> IO ThreadId
forkOn (I# cap) (IO action) =
  IO \s0 ->
    case forkOn# cap action s0 of
      (# s1, tid #) -> (# s1, ThreadId tid #)