{-# LANGUAGE
  BangPatterns,
  MagicHash,
  RankNTypes,
  ScopedTypeVariables,
  StandaloneKindSignatures,
  TypeOperators,
  UnboxedTuples #-}

-- | = Delimited continuations
--
-- Native multi-prompt delimited continuations.
-- These primitives let us manipulate slices of the call stack/evaluation
-- context delimited by 'reset'.
--
-- This module serves as a foundation for algebraic effect handlers,
-- a more structured interface for manipulating continuations and implementing
-- user-defined effects.
--
-- The behavior of 'reset' and 'shift0' is summarized by the following equations:
--
-- @
-- 'reset' (\\_ -> 'pure' x) = 'pure' x
-- 'reset' (\\t -> C ('shift0' t f)) = f (\\x -> 'reset' (\\t -> C x))
-- @
--
-- where @C@ is an evaluation context (in which @t@ may occur), i.e.,
-- a term of the following form:
--
-- > C x ::= C x >>= k      -- for any function k
-- >       | H (\h -> C x)  -- for any handler H ∈ { reset, (`runState` s), ... }
-- >       | x
--
--
-- This module ensures type safety. The rank-2 type of 'reset'
-- guarantees that 'shift0' will always have a maching 'reset' on the stack.
--
-- === References
--
-- - <https://ghc-proposals.readthedocs.io/en/latest/proposals/0313-delimited-continuation-primops.html Delimited continuation primops> (GHC proposal, implemented in GHC 9.6.1).
-- - <https://homes.luddy.indiana.edu/ccshan/recur/recur.pdf Shift to Control> (2004) by Chung-chieh Shan. The name 'shift0' follows the nomenclature in that paper.
module Bluefin.Algae.DelCont
  ( PromptTag
  , reset
  , shift0
  , Continuation
  , weakenC1
  , resume
  , continue
  , cancel
  ) where

import Data.Coerce (coerce)
import Data.Functor (void)
import Data.Kind (Type)
import GHC.Exts (State#, RealWorld, PromptTag#, prompt#, control0#, newPromptTag#)
import GHC.IO (IO(IO))
import Bluefin.Internal (Eff(UnsafeMkEff))
import Bluefin.Eff
import qualified Bluefin.Exception as E

-- | Tag for a prompt of type @Eff ss a@ and scope @s@.
type PromptTag :: Effects -> Type -> Effects -> Type
data PromptTag ss a s = MkPromptTag (PromptTag# a)

-- | Run the enclosed computation under a prompt of type @Eff ss a@.
--
-- @
-- f : forall s. 'PromptTag' ss a s -> 'Eff' (s ':&' ss) a
-- -------------------------------------------------
-- 'reset' (\\t -> f t) : 'Eff' ss a
-- @
--
-- The enclosed computation @f@ is given a tag which identifies that prompt
-- and remembers its type.
-- The scope parameter @s@ prevents the tag from being used outside of the
-- computation.
--
-- A prompt ('reset') delimits a slice of the call stack (or evaluation context),
-- which can be captured with 'shift0'. This slice, a continuation,
-- becomes a function of type @Eff ss0 b -> Eff ss a@ (where @Eff ss0 b@ is the
-- result type of 'shift0' at its calling site).
-- Calling the continuation restores the slice on the stack.
reset :: forall a ss.
  (forall s. PromptTag ss a s -> Eff (s :& ss) a) ->
  Eff ss a
reset :: forall a (ss :: Effects).
(forall (s :: Effects). PromptTag ss a s -> Eff (s :& ss) a)
-> Eff ss a
reset forall (s :: Effects). PromptTag ss a s -> Eff (s :& ss) a
f = IO# a -> Eff ss a
forall a (ss :: Effects). IO# a -> Eff ss a
unsafeMkEff (\State# RealWorld
z0 -> case State# RealWorld -> (# State# RealWorld, PromptTag# a #)
forall a. State# RealWorld -> (# State# RealWorld, PromptTag# a #)
newPromptTag# State# RealWorld
z0 of
    (# State# RealWorld
z1, PromptTag# a
tag #) -> PromptTag# a -> IO# a -> IO# a
forall a.
PromptTag# a
-> (State# RealWorld -> (# State# RealWorld, a #))
-> State# RealWorld
-> (# State# RealWorld, a #)
prompt# PromptTag# a
tag (Eff (Any :& ss) a -> IO# a
forall (ss :: Effects) a. Eff ss a -> IO# a
unsafeRunEff (PromptTag ss a Any -> Eff (Any :& ss) a
forall (s :: Effects). PromptTag ss a s -> Eff (s :& ss) a
f (PromptTag# a -> PromptTag ss a Any
forall (ss :: Effects) a (s :: Effects).
PromptTag# a -> PromptTag ss a s
MkPromptTag PromptTag# a
tag))) State# RealWorld
z1)

-- | Continuations are slices of the call stack, or evaluation context.
--
-- The 'Continuation' type is abstract since not all functions @'Eff' t b -> 'Eff' s a@
-- represent evaluation contexts. In particular, 'weakenC1' is not definable for arbitrary
-- such functions.
--
-- For example, in
--
-- @
-- reset \\tag0 ->
--   reset \\tag1 ->
--     reset \\tag2 ->
--       shift0 tag1 f >>= etc
-- @
--
-- 'shift0' captures a continuation, the slice represented by the following
-- function:
--
-- @
-- MkContinuation \\hole ->
--   reset \\tag1 ->
--     reset \\tag2 ->
--       hole >>= etc
-- @
--
-- That continuation has type @'Continuation' t s b a@ where @Eff t b@ is the type of the hole,
-- and @Eff s a@ is the type of the result once the hole is plugged.
--
-- The second argument of 'shift0', @f@, is applied to the continuation:
--
-- @
-- reset \\tag0 ->
--  f (MkContinuation \\hole ->
--   reset \\tag1 ->
--     reset \\tag2 ->
--       hole >>= etc)
-- @
newtype Continuation t s b a = MkContinuation (Eff t b -> Eff s a)

-- | Extend the context of a continuation.
weakenC1 :: Continuation t s b a -> Continuation (e :& t) (e :& s) b a
weakenC1 :: forall (t :: Effects) (s :: Effects) b a (e :: Effects).
Continuation t s b a -> Continuation (e :& t) (e :& s) b a
weakenC1 = Continuation t s b a -> Continuation (e :& t) (e :& s) b a
forall a b. Coercible a b => a -> b
coerce

-- | Resume a continuation with a computation under it.
resume :: Continuation t s b a -> Eff t b -> Eff s a
resume :: forall (t :: Effects) (s :: Effects) b a.
Continuation t s b a -> Eff t b -> Eff s a
resume (MkContinuation Eff t b -> Eff s a
k) = Eff t b -> Eff s a
k

-- | Resume a cancellable continuation with a result.
--
-- In other words, this converts a cancellable continuation to a simple continuation.
continue :: Continuation t s b a -> b -> Eff s a
continue :: forall (t :: Effects) (s :: Effects) b a.
Continuation t s b a -> b -> Eff s a
continue Continuation t s b a
k = Continuation t s b a -> Eff t b -> Eff s a
forall (t :: Effects) (s :: Effects) b a.
Continuation t s b a -> Eff t b -> Eff s a
resume Continuation t s b a
k (Eff t b -> Eff s a) -> (b -> Eff t b) -> b -> Eff s a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Eff t b
forall a. a -> Eff t a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

-- | Cancel a continuation: resume by throwing a scoped exception and catch it.
--
-- The continuation SHOULD re-throw unknown exceptions.
-- (That is guaranteed if you don't use "Exception.Dynamic".)
cancel :: Continuation t s b a -> Eff s ()
cancel :: forall (t :: Effects) (s :: Effects) b a.
Continuation t s b a -> Eff s ()
cancel Continuation t s b a
k = (forall (ex :: Effects). Exception () ex -> Eff (ex :& s) ())
-> (() -> Eff s ()) -> Eff s ()
forall e (es :: Effects) a.
(forall (ex :: Effects). Exception e ex -> Eff (ex :& es) a)
-> (e -> Eff es a) -> Eff es a
E.catch (\Exception () ex
ex -> Eff (ex :& s) a -> Eff (ex :& s) ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Continuation (ex :& t) (ex :& s) b a
-> Eff (ex :& t) b -> Eff (ex :& s) a
forall (t :: Effects) (s :: Effects) b a.
Continuation t s b a -> Eff t b -> Eff s a
resume (Continuation t s b a -> Continuation (ex :& t) (ex :& s) b a
forall (t :: Effects) (s :: Effects) b a (e :: Effects).
Continuation t s b a -> Continuation (e :& t) (e :& s) b a
weakenC1 Continuation t s b a
k) (Exception () ex -> () -> Eff (ex :& t) b
forall (ex :: Effects) (es :: Effects) e a.
(ex :> es) =>
Exception e ex -> e -> Eff es a
E.throw Exception () ex
ex ()))) (\()
_ -> () -> Eff s ()
forall a. a -> Eff s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())

-- | Capture the continuation up to the tagged prompt.
--
-- @
-- _ : s :> ss0
-- t : 'PromptTag' ss a s
-- f : ('Eff' ss0 b -> 'Eff' ss a) -> 'Eff' ss a
-- ---------------------------------------
-- 'shift0' t (\\k -> f k) : 'Eff' ss0 b
-- @
--
-- The prompt ('reset') is reinserted on the stack when the continuation is called:
--
-- @
-- 'reset' (\\t -> C ('shift0' t f)) = f (\\x -> 'reset' (\\t -> C x))
-- @
shift0 :: forall s a b ss ss0.
  s :> ss0 =>
  PromptTag ss a s ->
  (Continuation ss0 ss b a -> Eff ss a) ->
  Eff ss0 b
shift0 :: forall (s :: Effects) a b (ss :: Effects) (ss0 :: Effects).
(s :> ss0) =>
PromptTag ss a s
-> (Continuation ss0 ss b a -> Eff ss a) -> Eff ss0 b
shift0 (MkPromptTag PromptTag# a
tag) Continuation ss0 ss b a -> Eff ss a
f = IO# b -> Eff ss0 b
forall a (ss :: Effects). IO# a -> Eff ss a
unsafeMkEff (\State# RealWorld
z0 ->
  PromptTag# a
-> ((IO# b -> State# RealWorld -> (# State# RealWorld, a #))
    -> State# RealWorld -> (# State# RealWorld, a #))
-> IO# b
forall a b.
PromptTag# a
-> (((State# RealWorld -> (# State# RealWorld, b #))
     -> State# RealWorld -> (# State# RealWorld, a #))
    -> State# RealWorld -> (# State# RealWorld, a #))
-> State# RealWorld
-> (# State# RealWorld, b #)
control0# PromptTag# a
tag (\IO# b -> State# RealWorld -> (# State# RealWorld, a #)
k# ->
    Eff ss a -> State# RealWorld -> (# State# RealWorld, a #)
forall (ss :: Effects) a. Eff ss a -> IO# a
unsafeRunEff (Continuation ss0 ss b a -> Eff ss a
f ((IO# b -> State# RealWorld -> (# State# RealWorld, a #))
-> Continuation ss0 ss b a
forall b a (t :: Effects) (s :: Effects).
Continuation# b a -> Continuation t s b a
unsafeContinuation# (PromptTag# a
-> (State# RealWorld -> (# State# RealWorld, a #))
-> State# RealWorld
-> (# State# RealWorld, a #)
forall a.
PromptTag# a
-> (State# RealWorld -> (# State# RealWorld, a #))
-> State# RealWorld
-> (# State# RealWorld, a #)
prompt# PromptTag# a
tag ((State# RealWorld -> (# State# RealWorld, a #))
 -> State# RealWorld -> (# State# RealWorld, a #))
-> (IO# b -> State# RealWorld -> (# State# RealWorld, a #))
-> IO# b
-> State# RealWorld
-> (# State# RealWorld, a #)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO# b -> State# RealWorld -> (# State# RealWorld, a #)
k#)))) State# RealWorld
z0)

-- * Internal

type IO# a = State# RealWorld -> (# State# RealWorld , a #)
type Continuation# a b = IO# a -> IO# b

unsafeMkEff :: IO# a -> Eff ss a
unsafeMkEff :: forall a (ss :: Effects). IO# a -> Eff ss a
unsafeMkEff IO# a
f = IO a -> Eff ss a
forall (es :: Effects) a. IO a -> Eff es a
UnsafeMkEff (IO# a -> IO a
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO IO# a
f)

unsafeRunEff :: Eff ss a -> IO# a
unsafeRunEff :: forall (ss :: Effects) a. Eff ss a -> IO# a
unsafeRunEff (UnsafeMkEff (IO State# RealWorld -> (# State# RealWorld, a #)
f)) = State# RealWorld -> (# State# RealWorld, a #)
f

unsafeContinuation# :: Continuation# b a -> Continuation t s b a
unsafeContinuation# :: forall b a (t :: Effects) (s :: Effects).
Continuation# b a -> Continuation t s b a
unsafeContinuation# Continuation# b a
k = (Eff t b -> Eff s a) -> Continuation t s b a
forall (t :: Effects) (s :: Effects) b a.
(Eff t b -> Eff s a) -> Continuation t s b a
MkContinuation (IO# a -> Eff s a
forall a (ss :: Effects). IO# a -> Eff ss a
unsafeMkEff (IO# a -> Eff s a) -> (Eff t b -> IO# a) -> Eff t b -> Eff s a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Continuation# b a
k Continuation# b a -> (Eff t b -> IO# b) -> Eff t b -> IO# a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Eff t b -> IO# b
forall (ss :: Effects) a. Eff ss a -> IO# a
unsafeRunEff)