{-# OPTIONS_GHC -Wno-unused-foralls #-}

module Plutarch.Internal.TermCont (
  hashOpenTerm,
  TermCont (TermCont),
  runTermCont,
  unTermCont,
  tcont,
  pfindPlaceholder,
  pfindAllPlaceholders,
) where

import Data.Kind (Type)
import Data.List (nub)
import Data.String (fromString)
import Plutarch.Internal.Term (
  Config (Tracing),
  Dig,
  HoistedTerm (..),
  PType,
  RawTerm (..),
  S,
  Term (Term),
  TracingMode (DetTracing),
  asRawTerm,
  getTerm,
  hashRawTerm,
  perror,
  pgetConfig,
 )
import Plutarch.Internal.Trace (ptraceInfo)

newtype TermCont :: forall (r :: PType). S -> Type -> Type where
  TermCont :: forall r s a. {forall (r :: PType) (s :: S) a.
TermCont @r s a -> (a -> Term s r) -> Term s r
runTermCont :: (a -> Term s r) -> Term s r} -> TermCont @r s a

unTermCont :: TermCont @a s (Term s a) -> Term s a
unTermCont :: forall (a :: PType) (s :: S). TermCont @a s (Term s a) -> Term s a
unTermCont TermCont @a s (Term s a)
t = TermCont @a s (Term s a) -> (Term s a -> Term s a) -> Term s a
forall (r :: PType) (s :: S) a.
TermCont @r s a -> (a -> Term s r) -> Term s r
runTermCont TermCont @a s (Term s a)
t Term s a -> Term s a
forall a. a -> a
id

instance Functor (TermCont s) where
  fmap :: forall a b. (a -> b) -> TermCont @r s a -> TermCont @r s b
fmap a -> b
f (TermCont (a -> Term s r) -> Term s r
g) = ((b -> Term s r) -> Term s r) -> TermCont @r s b
forall (r :: PType) (s :: S) a.
((a -> Term s r) -> Term s r) -> TermCont @r s a
TermCont (((b -> Term s r) -> Term s r) -> TermCont @r s b)
-> ((b -> Term s r) -> Term s r) -> TermCont @r s b
forall a b. (a -> b) -> a -> b
$ \b -> Term s r
h -> (a -> Term s r) -> Term s r
g (b -> Term s r
h (b -> Term s r) -> (a -> b) -> a -> Term s r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f)

instance Applicative (TermCont s) where
  pure :: forall a. a -> TermCont @r s a
pure a
x = ((a -> Term s r) -> Term s r) -> TermCont @r s a
forall (r :: PType) (s :: S) a.
((a -> Term s r) -> Term s r) -> TermCont @r s a
TermCont (((a -> Term s r) -> Term s r) -> TermCont @r s a)
-> ((a -> Term s r) -> Term s r) -> TermCont @r s a
forall a b. (a -> b) -> a -> b
$ \a -> Term s r
f -> a -> Term s r
f a
x
  TermCont @r s (a -> b)
x <*> :: forall a b.
TermCont @r s (a -> b) -> TermCont @r s a -> TermCont @r s b
<*> TermCont @r s a
y = do
    a -> b
x <- TermCont @r s (a -> b)
x
    a -> b
x (a -> b) -> TermCont @r s a -> TermCont @r s b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> TermCont @r s a
y

instance Monad (TermCont s) where
  (TermCont (a -> Term s r) -> Term s r
f) >>= :: forall a b.
TermCont @r s a -> (a -> TermCont @r s b) -> TermCont @r s b
>>= a -> TermCont @r s b
g = ((b -> Term s r) -> Term s r) -> TermCont @r s b
forall (r :: PType) (s :: S) a.
((a -> Term s r) -> Term s r) -> TermCont @r s a
TermCont (((b -> Term s r) -> Term s r) -> TermCont @r s b)
-> ((b -> Term s r) -> Term s r) -> TermCont @r s b
forall a b. (a -> b) -> a -> b
$ \b -> Term s r
h ->
    (a -> Term s r) -> Term s r
f
      ( \a
x ->
          TermCont @r s b -> (b -> Term s r) -> Term s r
forall (r :: PType) (s :: S) a.
TermCont @r s a -> (a -> Term s r) -> Term s r
runTermCont (a -> TermCont @r s b
g a
x) b -> Term s r
h
      )

instance MonadFail (TermCont s) where
  fail :: forall a. String -> TermCont @r s a
fail String
s = ((a -> Term s r) -> Term s r) -> TermCont @r s a
forall (r :: PType) (s :: S) a.
((a -> Term s r) -> Term s r) -> TermCont @r s a
TermCont (((a -> Term s r) -> Term s r) -> TermCont @r s a)
-> ((a -> Term s r) -> Term s r) -> TermCont @r s a
forall a b. (a -> b) -> a -> b
$ \a -> Term s r
_ ->
    (Config -> Term s r) -> Term s r
forall (s :: S) (a :: PType). (Config -> Term s a) -> Term s a
pgetConfig ((Config -> Term s r) -> Term s r)
-> (Config -> Term s r) -> Term s r
forall a b. (a -> b) -> a -> b
$ \case
      -- Note: This currently works because DetTracing is the most specific
      -- tracing mode.
      Tracing LogLevel
_ TracingMode
DetTracing -> 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
"Pattern matching failure in TermCont" Term s r
forall (s :: S) (a :: PType). Term s a
perror
      Config
_ -> Term s PString -> Term s r -> Term s r
forall (a :: PType) (s :: S).
Term s PString -> Term s a -> Term s a
ptraceInfo (String -> Term s PString
forall a. IsString a => String -> a
fromString String
s) Term s r
forall (s :: S) (a :: PType). Term s a
perror

tcont :: ((a -> Term s r) -> Term s r) -> TermCont @r s a
tcont :: forall a (s :: S) (r :: PType).
((a -> Term s r) -> Term s r) -> TermCont @r s a
tcont = ((a -> Term s r) -> Term s r) -> TermCont @r s a
forall (r :: PType) (s :: S) a.
((a -> Term s r) -> Term s r) -> TermCont @r s a
TermCont

hashOpenTerm :: Term s a -> TermCont s Dig
hashOpenTerm :: forall {r :: PType} (s :: S) (a :: PType).
Term s a -> TermCont @r s Dig
hashOpenTerm Term s a
x = ((Dig -> Term s r) -> Term s r) -> TermCont @r s Dig
forall (r :: PType) (s :: S) a.
((a -> Term s r) -> Term s r) -> TermCont @r s a
TermCont (((Dig -> Term s r) -> Term s r) -> TermCont @r s Dig)
-> ((Dig -> Term s r) -> Term s r) -> TermCont @r s Dig
forall a b. (a -> b) -> a -> b
$ \Dig -> Term s r
f -> (Word64 -> TermMonad TermResult) -> Term s r
forall (s :: S) (a :: PType).
(Word64 -> TermMonad TermResult) -> Term s a
Term ((Word64 -> TermMonad TermResult) -> Term s r)
-> (Word64 -> TermMonad TermResult) -> Term s r
forall a b. (a -> b) -> a -> b
$ \Word64
i -> do
  TermResult
y <- Term s a -> Word64 -> TermMonad TermResult
forall (s :: S) (a :: PType).
Term s a -> Word64 -> TermMonad TermResult
asRawTerm Term s a
x Word64
i
  Term s r -> Word64 -> TermMonad TermResult
forall (s :: S) (a :: PType).
Term s a -> Word64 -> TermMonad TermResult
asRawTerm (Dig -> Term s r
f (Dig -> Term s r) -> (TermResult -> Dig) -> TermResult -> Term s r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RawTerm -> Dig
hashRawTerm (RawTerm -> Dig) -> (TermResult -> RawTerm) -> TermResult -> Dig
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TermResult -> RawTerm
getTerm (TermResult -> Term s r) -> TermResult -> Term s r
forall a b. (a -> b) -> a -> b
$ TermResult
y) Word64
i

-- This can technically be done outside of TermCont.
-- Need to pay close attention when killing branch with this.
-- If term is pre-evaluated (via `evalTerm`), RawTerm will no longer hold
-- tagged RPlaceholder.

{- | Given a term, and an integer tag, this function checks if the term holds and
@PPlaceholder@ with the given integer tag.
-}
pfindPlaceholder :: Integer -> Term s a -> TermCont s Bool
pfindPlaceholder :: forall {r :: PType} (s :: S) (a :: PType).
Integer -> Term s a -> TermCont @r s Bool
pfindPlaceholder Integer
idx Term s a
x = ((Bool -> Term s r) -> Term s r) -> TermCont @r s Bool
forall (r :: PType) (s :: S) a.
((a -> Term s r) -> Term s r) -> TermCont @r s a
TermCont (((Bool -> Term s r) -> Term s r) -> TermCont @r s Bool)
-> ((Bool -> Term s r) -> Term s r) -> TermCont @r s Bool
forall a b. (a -> b) -> a -> b
$ \Bool -> Term s r
f -> (Word64 -> TermMonad TermResult) -> Term s r
forall (s :: S) (a :: PType).
(Word64 -> TermMonad TermResult) -> Term s a
Term ((Word64 -> TermMonad TermResult) -> Term s r)
-> (Word64 -> TermMonad TermResult) -> Term s r
forall a b. (a -> b) -> a -> b
$ \Word64
i -> do
  TermResult
y <- Term s a -> Word64 -> TermMonad TermResult
forall (s :: S) (a :: PType).
Term s a -> Word64 -> TermMonad TermResult
asRawTerm Term s a
x Word64
i

  let
    findPlaceholder :: RawTerm -> Bool
findPlaceholder (RLamAbs Word64
_ RawTerm
x) = RawTerm -> Bool
findPlaceholder RawTerm
x
    findPlaceholder (RApply RawTerm
x [RawTerm]
xs) = (RawTerm -> Bool) -> [RawTerm] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any RawTerm -> Bool
findPlaceholder (RawTerm
x RawTerm -> [RawTerm] -> [RawTerm]
forall a. a -> [a] -> [a]
: [RawTerm]
xs)
    findPlaceholder (RForce RawTerm
x) = RawTerm -> Bool
findPlaceholder RawTerm
x
    findPlaceholder (RDelay RawTerm
x) = RawTerm -> Bool
findPlaceholder RawTerm
x
    findPlaceholder (RHoisted (HoistedTerm Dig
_ RawTerm
x)) = RawTerm -> Bool
findPlaceholder RawTerm
x
    findPlaceholder (RPlaceHolder Integer
idx') = Integer
idx Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
idx'
    findPlaceholder (RConstr Word64
_ [RawTerm]
xs) = (RawTerm -> Bool) -> [RawTerm] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any RawTerm -> Bool
findPlaceholder [RawTerm]
xs
    findPlaceholder (RCase RawTerm
x [RawTerm]
xs) = (RawTerm -> Bool) -> [RawTerm] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any RawTerm -> Bool
findPlaceholder (RawTerm
x RawTerm -> [RawTerm] -> [RawTerm]
forall a. a -> [a] -> [a]
: [RawTerm]
xs)
    findPlaceholder (RVar Word64
_) = Bool
False
    findPlaceholder (RConstant Some @Type (ValueOf DefaultUni)
_) = Bool
False
    findPlaceholder (RBuiltin DefaultFun
_) = Bool
False
    findPlaceholder (RCompiled UTerm
_) = Bool
False
    findPlaceholder RawTerm
RError = Bool
False

  Term s r -> Word64 -> TermMonad TermResult
forall (s :: S) (a :: PType).
Term s a -> Word64 -> TermMonad TermResult
asRawTerm (Bool -> Term s r
f (Bool -> Term s r)
-> (TermResult -> Bool) -> TermResult -> Term s r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RawTerm -> Bool
findPlaceholder (RawTerm -> Bool) -> (TermResult -> RawTerm) -> TermResult -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TermResult -> RawTerm
getTerm (TermResult -> Term s r) -> TermResult -> Term s r
forall a b. (a -> b) -> a -> b
$ TermResult
y) Word64
i

-- | Finds all placeholder ids and returns it
pfindAllPlaceholders :: Term s a -> TermCont s [Integer]
pfindAllPlaceholders :: forall {r :: PType} (s :: S) (a :: PType).
Term s a -> TermCont @r s [Integer]
pfindAllPlaceholders Term s a
x = (([Integer] -> Term s r) -> Term s r) -> TermCont @r s [Integer]
forall (r :: PType) (s :: S) a.
((a -> Term s r) -> Term s r) -> TermCont @r s a
TermCont ((([Integer] -> Term s r) -> Term s r) -> TermCont @r s [Integer])
-> (([Integer] -> Term s r) -> Term s r) -> TermCont @r s [Integer]
forall a b. (a -> b) -> a -> b
$ \[Integer] -> Term s r
f -> (Word64 -> TermMonad TermResult) -> Term s r
forall (s :: S) (a :: PType).
(Word64 -> TermMonad TermResult) -> Term s a
Term ((Word64 -> TermMonad TermResult) -> Term s r)
-> (Word64 -> TermMonad TermResult) -> Term s r
forall a b. (a -> b) -> a -> b
$ \Word64
i -> do
  TermResult
y <- Term s a -> Word64 -> TermMonad TermResult
forall (s :: S) (a :: PType).
Term s a -> Word64 -> TermMonad TermResult
asRawTerm Term s a
x Word64
i

  let
    findPlaceholder :: RawTerm -> t
findPlaceholder (RLamAbs Word64
_ RawTerm
x) = RawTerm -> t
findPlaceholder RawTerm
x
    findPlaceholder (RApply RawTerm
x [RawTerm]
xs) = RawTerm -> t
findPlaceholder RawTerm
x t -> t -> t
forall a. Semigroup a => a -> a -> a
<> (RawTerm -> t) -> [RawTerm] -> t
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: Type -> Type) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap RawTerm -> t
findPlaceholder [RawTerm]
xs
    findPlaceholder (RForce RawTerm
x) = RawTerm -> t
findPlaceholder RawTerm
x
    findPlaceholder (RDelay RawTerm
x) = RawTerm -> t
findPlaceholder RawTerm
x
    findPlaceholder (RHoisted (HoistedTerm Dig
_ RawTerm
x)) = RawTerm -> t
findPlaceholder RawTerm
x
    findPlaceholder (RPlaceHolder Integer
idx) = [Integer
Item t
idx]
    findPlaceholder (RConstr Word64
_ [RawTerm]
xs) = (RawTerm -> t) -> [RawTerm] -> t
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: Type -> Type) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap RawTerm -> t
findPlaceholder [RawTerm]
xs
    findPlaceholder (RCase RawTerm
x [RawTerm]
xs) = RawTerm -> t
findPlaceholder RawTerm
x t -> t -> t
forall a. Semigroup a => a -> a -> a
<> (RawTerm -> t) -> [RawTerm] -> t
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: Type -> Type) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap RawTerm -> t
findPlaceholder [RawTerm]
xs
    findPlaceholder (RVar Word64
_) = []
    findPlaceholder (RConstant Some @Type (ValueOf DefaultUni)
_) = []
    findPlaceholder (RBuiltin DefaultFun
_) = []
    findPlaceholder (RCompiled UTerm
_) = []
    findPlaceholder RawTerm
RError = []

  Term s r -> Word64 -> TermMonad TermResult
forall (s :: S) (a :: PType).
Term s a -> Word64 -> TermMonad TermResult
asRawTerm ([Integer] -> Term s r
f ([Integer] -> Term s r)
-> (TermResult -> [Integer]) -> TermResult -> Term s r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Integer] -> [Integer]
forall a. Eq a => [a] -> [a]
nub ([Integer] -> [Integer])
-> (TermResult -> [Integer]) -> TermResult -> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RawTerm -> [Integer]
forall {t}.
((Item t :: Type) ~ (Integer :: Type), Monoid t, IsList t) =>
RawTerm -> t
findPlaceholder (RawTerm -> [Integer])
-> (TermResult -> RawTerm) -> TermResult -> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TermResult -> RawTerm
getTerm (TermResult -> Term s r) -> TermResult -> Term s r
forall a b. (a -> b) -> a -> b
$ TermResult
y) Word64
i