{-# LANGUAGE AllowAmbiguousTypes #-}

module Plutarch.TermCont (
  TC.hashOpenTerm,
  TC.TermCont (TermCont),
  TC.runTermCont,
  TC.unTermCont,
  TC.tcont,
  pletC,
  pmatchC,
  pletFieldsC,
  ptraceC,
  pguardC,
  pguardC',
  ptryFromC,
  pexpectJustC,
) where

import Data.Kind (Type)
import Plutarch.Builtin.Bool (PBool, pif)
import Plutarch.Builtin.String (PString)
import Plutarch.DataRepr (HRec, PDataFields, PFields, pletFields)
import Plutarch.DataRepr.Internal.Field (
  BindFields,
  Bindings,
  BoundTerms,
 )
import Plutarch.Internal.PlutusType (PlutusType, pmatch)
import Plutarch.Internal.Term (S, Term, plet)
import Plutarch.Internal.TermCont (TermCont, tcont)
import Plutarch.Internal.TermCont qualified as TC
import Plutarch.Internal.TryFrom (PTryFrom (PTryFromExcess), ptryFrom)
import Plutarch.Maybe (PMaybe (PJust, PNothing))
import Plutarch.Reducible (Reduce)
import Plutarch.Trace (ptraceInfo, ptraceInfoError)

-- | Like `plet` but works in a `TermCont` monad
pletC :: Term s a -> TermCont s (Term s a)
pletC :: forall {r :: PType} (s :: S) (a :: PType).
Term s a -> TermCont @r s (Term s a)
pletC = ((Term s a -> Term s r) -> Term s r) -> TermCont @r s (Term s a)
forall a (s :: S) (r :: PType).
((a -> Term s r) -> Term s r) -> TermCont @r s a
tcont (((Term s a -> Term s r) -> Term s r) -> TermCont @r s (Term s a))
-> (Term s a -> (Term s a -> Term s r) -> Term s r)
-> Term s a
-> TermCont @r s (Term s a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term s a -> (Term s a -> Term s r) -> Term s r
forall (s :: S) (a :: PType) (b :: PType).
Term s a -> (Term s a -> Term s b) -> Term s b
plet

-- | Like `pmatch` but works in a `TermCont` monad
pmatchC :: PlutusType a => Term s a -> TermCont s (a s)
pmatchC :: forall {r :: PType} (a :: PType) (s :: S).
PlutusType a =>
Term s a -> TermCont @r s (a s)
pmatchC = ((a s -> Term s r) -> Term s r) -> TermCont @r s (a s)
forall a (s :: S) (r :: PType).
((a -> Term s r) -> Term s r) -> TermCont @r s a
tcont (((a s -> Term s r) -> Term s r) -> TermCont @r s (a s))
-> (Term s a -> (a s -> Term s r) -> Term s r)
-> Term s a
-> TermCont @r s (a s)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term s a -> (a s -> Term s r) -> Term s r
forall (a :: PType) (s :: S) (b :: PType).
PlutusType a =>
Term s a -> (a s -> Term s b) -> Term s b
pmatch

-- | Like `pletFields` but works in a `TermCont` monad.
pletFieldsC ::
  forall fs a s b ps bs.
  ( PDataFields a
  , ps ~ PFields a
  , bs ~ Bindings ps fs
  , BindFields ps bs
  ) =>
  Term s a ->
  TermCont @b s (HRec (BoundTerms ps bs s))
pletFieldsC :: forall (fs :: [Symbol]) (a :: PType) (s :: S) (b :: PType)
       (ps :: [PLabeledType]) (bs :: [ToBind]).
(PDataFields a,
 (ps :: [PLabeledType]) ~ (PFields a :: [PLabeledType]),
 (bs :: [ToBind]) ~ (Bindings ps fs :: [ToBind]),
 BindFields ps bs) =>
Term s a -> TermCont @b s (HRec (BoundTerms ps bs s))
pletFieldsC Term s a
x = ((HRec (BoundTerms ps bs s) -> Term s b) -> Term s b)
-> TermCont @b s (HRec (BoundTerms ps bs s))
forall a (s :: S) (r :: PType).
((a -> Term s r) -> Term s r) -> TermCont @r s a
tcont (((HRec (BoundTerms ps bs s) -> Term s b) -> Term s b)
 -> TermCont @b s (HRec (BoundTerms ps bs s)))
-> ((HRec (BoundTerms ps bs s) -> Term s b) -> Term s b)
-> TermCont @b s (HRec (BoundTerms ps bs s))
forall a b. (a -> b) -> a -> b
$ forall (fs :: [Symbol]) (a :: PType) (s :: S) (b :: PType)
       (ps :: [PLabeledType]) (bs :: [ToBind]).
(PDataFields a,
 (ps :: [PLabeledType]) ~ (PFields a :: [PLabeledType]),
 (bs :: [ToBind]) ~ (Bindings ps fs :: [ToBind]),
 BindFields ps bs) =>
Term s a -> (HRecOf a fs s -> Term s b) -> Term s b
pletFields @fs Term s a
x

{- | Like `ptrace` but works in a `TermCont` monad.

=== Example ===

@
foo :: Term s PUnit
foo = unTermCont $ do
  ptraceC "returning unit!"
  pure $ pconstant ()
@
-}
ptraceC :: Term s PString -> TermCont s ()
ptraceC :: forall {r :: PType} (s :: S). Term s PString -> TermCont @r s ()
ptraceC Term s PString
s = ((() -> Term s r) -> Term s r) -> TermCont @r s ()
forall a (s :: S) (r :: PType).
((a -> Term s r) -> Term s r) -> TermCont @r s a
tcont (((() -> Term s r) -> Term s r) -> TermCont @r s ())
-> ((() -> Term s r) -> Term s r) -> TermCont @r s ()
forall a b. (a -> b) -> a -> b
$ \() -> Term s r
f -> Term s PString -> Term s r -> Term s r
forall (a :: PType) (s :: S).
Term s PString -> Term s a -> Term s a
ptraceInfo Term s PString
s (() -> Term s r
f ())

{- | Trace a message and raise error if 'cond' is false. Otherwise, continue.

=== Example ===

@
onlyAllow42 :: Term s (PInteger :--> PUnit)
onlyAllow42 = plam $ \i -> unTermCont $ do
  pguardC "expected 42" $ i #== 42
  pure $ pconstant ()
@
-}
pguardC :: Term s PString -> Term s PBool -> TermCont s ()
pguardC :: forall {r :: PType} (s :: S).
Term s PString -> Term s PBool -> TermCont @r s ()
pguardC Term s PString
s Term s PBool
cond = ((() -> Term s r) -> Term s r) -> TermCont @r s ()
forall a (s :: S) (r :: PType).
((a -> Term s r) -> Term s r) -> TermCont @r s a
tcont (((() -> Term s r) -> Term s r) -> TermCont @r s ())
-> ((() -> Term s r) -> Term s r) -> TermCont @r s ()
forall a b. (a -> b) -> a -> b
$ \() -> Term s r
f -> Term s PBool -> Term s r -> Term s r -> Term s r
forall (a :: PType) (s :: S).
Term s PBool -> Term s a -> Term s a -> Term s a
pif Term s PBool
cond (() -> Term s r
f ()) (Term s r -> Term s r) -> Term s r -> Term s r
forall a b. (a -> b) -> a -> b
$ Term s PString -> Term s r
forall (a :: PType) (s :: S). Term s PString -> Term s a
ptraceInfoError Term s PString
s

{- | Stop computation and return given term if 'cond' is false. Otherwise, continue.

=== Example ===

@
is42 :: Term s (PInteger :--> PBool)
is42 = plam $ \i -> unTermCont $ do
  pguardC' (pconstant False) $ i #== 42
  pure $ pconstant True
@
-}
pguardC' :: Term s a -> Term s PBool -> TermCont @a s ()
pguardC' :: forall (s :: S) (a :: PType).
Term s a -> Term s PBool -> TermCont @a s ()
pguardC' Term s a
r Term s PBool
cond = ((() -> Term s a) -> Term s a) -> TermCont @a s ()
forall a (s :: S) (r :: PType).
((a -> Term s r) -> Term s r) -> TermCont @r s a
tcont (((() -> Term s a) -> Term s a) -> TermCont @a s ())
-> ((() -> Term s a) -> Term s a) -> TermCont @a s ()
forall a b. (a -> b) -> a -> b
$ \() -> Term s a
f -> Term s PBool -> Term s a -> Term s a -> Term s a
forall (a :: PType) (s :: S).
Term s PBool -> Term s a -> Term s a -> Term s a
pif Term s PBool
cond (() -> Term s a
f ()) Term s a
r

-- | 'TermCont' producing version of 'ptryFrom'.
ptryFromC :: forall b r a s. PTryFrom a b => Term s a -> TermCont @r s (Term s b, Reduce (PTryFromExcess a b s))
ptryFromC :: forall (b :: PType) (r :: PType) (a :: PType) (s :: S).
PTryFrom a b =>
Term s a -> TermCont @r s (Term s b, Reduce (PTryFromExcess a b s))
ptryFromC = (((Term s b, Reduce (PTryFromExcess a b s)) -> Term s r)
 -> Term s r)
-> TermCont @r s (Term s b, Reduce (PTryFromExcess a b s))
forall a (s :: S) (r :: PType).
((a -> Term s r) -> Term s r) -> TermCont @r s a
tcont ((((Term s b, Reduce (PTryFromExcess a b s)) -> Term s r)
  -> Term s r)
 -> TermCont @r s (Term s b, Reduce (PTryFromExcess a b s)))
-> (Term s a
    -> ((Term s b, Reduce (PTryFromExcess a b s)) -> Term s r)
    -> Term s r)
-> Term s a
-> TermCont @r s (Term s b, Reduce (PTryFromExcess a b s))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term s a
-> ((Term s b, Reduce (PTryFromExcess a b s)) -> Term s r)
-> Term s r
forall (b :: PType) (a :: PType) (s :: S) (r :: PType).
PTryFrom a b =>
Term s a
-> ((Term s b, Reduce (PTryFromExcess a b s)) -> Term s r)
-> Term s r
ptryFrom

{- | Escape with a particular value on expecting 'PJust'. For use in monadic context.

@since WIP
-}
pexpectJustC ::
  forall (a :: S -> Type) (r :: S -> Type) (s :: S).
  Term s r ->
  Term s (PMaybe a) ->
  TermCont @r s (Term s a)
pexpectJustC :: forall (a :: PType) (r :: PType) (s :: S).
Term s r -> Term s (PMaybe a) -> TermCont @r s (Term s a)
pexpectJustC Term s r
escape Term s (PMaybe a)
ma = ((Term s a -> Term s r) -> Term s r) -> TermCont @r s (Term s a)
forall a (s :: S) (r :: PType).
((a -> Term s r) -> Term s r) -> TermCont @r s a
tcont (((Term s a -> Term s r) -> Term s r) -> TermCont @r s (Term s a))
-> ((Term s a -> Term s r) -> Term s r) -> TermCont @r s (Term s a)
forall a b. (a -> b) -> a -> b
$ \Term s a -> Term s r
f ->
  Term s (PMaybe a) -> (PMaybe a s -> Term s r) -> Term s r
forall (a :: PType) (s :: S) (b :: PType).
PlutusType a =>
Term s a -> (a s -> Term s b) -> Term s b
pmatch Term s (PMaybe a)
ma ((PMaybe a s -> Term s r) -> Term s r)
-> (PMaybe a s -> Term s r) -> Term s r
forall a b. (a -> b) -> a -> b
$ \case
    PJust Term s a
v -> Term s a -> Term s r
f Term s a
v
    PMaybe a s
PNothing -> Term s r
escape