module UnconditionalJump
  ( -- * Labels
    Label,
    label,
    goto,

    -- ** Derived @label@ variants
    labelWith,
    labelE,
  )
where

import Control.Exception (Exception (..), SomeException, asyncExceptionFromException, asyncExceptionToException, catch, throwIO)
import Data.Functor.Contravariant (Contravariant (contramap))
import IntSupply (IntSupply)
import IntSupply qualified
import System.IO.Unsafe (unsafePerformIO)
import Unsafe.Coerce (unsafeCoerce)

newtype Label a
  = Label (forall x. a -> IO x)

instance Contravariant Label where
  contramap :: forall a' a. (a' -> a) -> Label a -> Label a'
contramap a' -> a
f (Label forall x. a -> IO x
g) =
    (forall x. a' -> IO x) -> Label a'
forall a. (forall x. a -> IO x) -> Label a
Label (a -> IO x
forall x. a -> IO x
g (a -> IO x) -> (a' -> a) -> a' -> IO x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a' -> a
f)

-- | Create a label.
label :: (Label a -> IO a) -> IO a
label :: forall a. (Label a -> IO a) -> IO a
label Label a -> IO a
f = do
  Int
i <- IntSupply -> IO Int
IntSupply.next IntSupply
supply
  IO a -> (X -> IO a) -> IO a
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
catch (Label a -> IO a
f ((forall x. a -> IO x) -> Label a
forall a. (forall x. a -> IO x) -> Label a
Label (X -> IO x
forall e a. Exception e => e -> IO a
throwIO (X -> IO x) -> (a -> X) -> a -> IO x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> a -> X
forall a. Int -> a -> X
X Int
i))) \err :: X
err@(X Int
j a
x) ->
    if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j
      then a -> IO a
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> a
forall a b. a -> b
unsafeCoerce a
x)
      else X -> IO a
forall e a. Exception e => e -> IO a
throwIO X
err

-- |
-- @
-- labelWith f g h = label (fmap g . h . contramap f)
-- @
labelWith :: (a -> c) -> (b -> c) -> (Label a -> IO b) -> IO c
labelWith :: forall a c b. (a -> c) -> (b -> c) -> (Label a -> IO b) -> IO c
labelWith a -> c
f b -> c
g Label a -> IO b
action =
  (Label c -> IO c) -> IO c
forall a. (Label a -> IO a) -> IO a
label ((b -> c) -> IO b -> IO c
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> c
g (IO b -> IO c) -> (Label c -> IO b) -> Label c -> IO c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Label a -> IO b
action (Label a -> IO b) -> (Label c -> Label a) -> Label c -> IO b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> c) -> Label c -> Label a
forall a' a. (a' -> a) -> Label a -> Label a'
forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap a -> c
f)

-- |
-- @
-- labelE = labelWith Left Right
-- @
labelE :: (Label a -> IO b) -> IO (Either a b)
labelE :: forall a b. (Label a -> IO b) -> IO (Either a b)
labelE =
  (a -> Either a b)
-> (b -> Either a b) -> (Label a -> IO b) -> IO (Either a b)
forall a c b. (a -> c) -> (b -> c) -> (Label a -> IO b) -> IO c
labelWith a -> Either a b
forall a b. a -> Either a b
Left b -> Either a b
forall a b. b -> Either a b
Right

-- | Go to a label.
goto :: Label a -> a -> IO notreached
goto :: forall a notreached. Label a -> a -> IO notreached
goto (Label forall x. a -> IO x
f) a
x =
  a -> IO notreached
forall x. a -> IO x
f a
x

data X = forall a. X {-# UNPACK #-} !Int a

-- Make X an async exception so it's less likely to be caught and ignored
instance Exception X where
  toException :: X -> SomeException
  toException :: X -> SomeException
toException = X -> SomeException
forall e. Exception e => e -> SomeException
asyncExceptionToException

  fromException :: SomeException -> Maybe X
  fromException :: SomeException -> Maybe X
fromException = SomeException -> Maybe X
forall e. Exception e => SomeException -> Maybe e
asyncExceptionFromException

instance Show X where
  show :: X -> String
show X
_ = String
"«unconditional-jump»"

supply :: IntSupply
supply :: IntSupply
supply = IO IntSupply -> IntSupply
forall a. IO a -> a
unsafePerformIO IO IntSupply
IntSupply.new
{-# NOINLINE supply #-}