-- | Utilities for unit testing plutarch terms
module Plutarch.Test.Unit (
  testCompileFail,
  testEval,
  testEvalFail,
  testEvalEqual,
  testEvalEqualTraces,
  TermResult (..),
  evalTermResult,
) where

import Data.Kind (Type)
import Data.Text (Text)
import Data.Text qualified as Text
import Plutarch.Evaluate (evalScriptUnlimited)
import Plutarch.Internal.Other (printScript)
import Plutarch.Internal.Term (
  Config (NoTracing, Tracing),
  LogLevel (LogDebug),
  TracingMode (DetTracing),
  compile,
 )
import Plutarch.Prelude
import PlutusCore qualified as PLC (DefaultFun, DefaultUni, NamedDeBruijn)
import Test.Tasty (TestName, TestTree)
import Test.Tasty.HUnit (assertEqual, assertFailure, testCase)
import UntypedPlutusCore.Evaluation.Machine.Cek qualified as Cek

{- | Assert that term compiled and evaluated without errors

@since 1.0.0
-}
testEval :: forall (a :: S -> Type). TestName -> (forall (s :: S). Term s a) -> TestTree
testEval :: forall (a :: S -> Type).
TestName -> (forall (s :: S). Term s a) -> TestTree
testEval TestName
name forall (s :: S). Term s a
term = TestName -> Assertion -> TestTree
testCase TestName
name (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ do
  case Config -> (forall (s :: S). Term s a) -> TermResult
forall (a :: S -> Type).
Config -> (forall (s :: S). Term s a) -> TermResult
evalTermResult (LogLevel -> TracingMode -> Config
Tracing LogLevel
LogDebug TracingMode
DetTracing) Term s a
forall (s :: S). Term s a
term of
    FailedToCompile Text
err -> TestName -> Assertion
forall a. HasCallStack => TestName -> IO a
assertFailure (TestName -> Assertion) -> TestName -> Assertion
forall a b. (a -> b) -> a -> b
$ TestName
"Failed to compile: " TestName -> TestName -> TestName
forall a. Semigroup a => a -> a -> a
<> Text -> TestName
Text.unpack Text
err
    FailedToEvaluate CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
err [Text]
_ -> TestName -> Assertion
forall a. HasCallStack => TestName -> IO a
assertFailure (TestName -> Assertion) -> TestName -> Assertion
forall a b. (a -> b) -> a -> b
$ TestName
"Failed to evaluate: " TestName -> TestName -> TestName
forall a. Semigroup a => a -> a -> a
<> CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
-> TestName
forall a. Show a => a -> TestName
show CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
err
    Evaluated TestName
_ [Text]
_ -> () -> Assertion
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ()

{- | Assert that term compiled correctly but evaluated with errors

@since 1.0.0
-}
testEvalFail :: forall (a :: S -> Type). TestName -> (forall (s :: S). Term s a) -> TestTree
testEvalFail :: forall (a :: S -> Type).
TestName -> (forall (s :: S). Term s a) -> TestTree
testEvalFail TestName
name forall (s :: S). Term s a
term = TestName -> Assertion -> TestTree
testCase TestName
name (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ do
  case Config -> (forall (s :: S). Term s a) -> TermResult
forall (a :: S -> Type).
Config -> (forall (s :: S). Term s a) -> TermResult
evalTermResult Config
NoTracing Term s a
forall (s :: S). Term s a
term of
    FailedToCompile Text
err -> TestName -> Assertion
forall a. HasCallStack => TestName -> IO a
assertFailure (TestName -> Assertion) -> TestName -> Assertion
forall a b. (a -> b) -> a -> b
$ TestName
"Failed to compile: " TestName -> TestName -> TestName
forall a. Semigroup a => a -> a -> a
<> Text -> TestName
Text.unpack Text
err
    FailedToEvaluate CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
_ [Text]
_ -> () -> Assertion
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ()
    Evaluated TestName
script [Text]
_ -> TestName -> Assertion
forall a. HasCallStack => TestName -> IO a
assertFailure (TestName
"Evaluated, but expected failure:\n" TestName -> TestName -> TestName
forall a. Semigroup a => a -> a -> a
<> TestName
script)

{- | Assert that term failed to compile

@since 1.0.0
-}
testCompileFail :: forall (a :: S -> Type). TestName -> (forall (s :: S). Term s a) -> TestTree
testCompileFail :: forall (a :: S -> Type).
TestName -> (forall (s :: S). Term s a) -> TestTree
testCompileFail TestName
name forall (s :: S). Term s a
term = TestName -> Assertion -> TestTree
testCase TestName
name (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ do
  case Config -> (forall (s :: S). Term s a) -> TermResult
forall (a :: S -> Type).
Config -> (forall (s :: S). Term s a) -> TermResult
evalTermResult Config
NoTracing Term s a
forall (s :: S). Term s a
term of
    FailedToCompile Text
_ -> () -> Assertion
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ()
    FailedToEvaluate CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
err [Text]
_ -> TestName -> Assertion
forall a. HasCallStack => TestName -> IO a
assertFailure (TestName -> Assertion) -> TestName -> Assertion
forall a b. (a -> b) -> a -> b
$ TestName
"Failed to evaluate: " TestName -> TestName -> TestName
forall a. Semigroup a => a -> a -> a
<> CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
-> TestName
forall a. Show a => a -> TestName
show CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
err
    Evaluated TestName
script [Text]
_ -> TestName -> Assertion
forall a. HasCallStack => TestName -> IO a
assertFailure (TestName -> Assertion) -> TestName -> Assertion
forall a b. (a -> b) -> a -> b
$ TestName
"Evaluated, but expected failure: " TestName -> TestName -> TestName
forall a. Semigroup a => a -> a -> a
<> TestName
script

{- | Assert that term compiled and evaluated without errors and matches the expected value
note that comparison is done on AST level, not by `Eq` or `PEq`

@since 1.0.0
-}
testEvalEqual ::
  forall (a :: S -> Type).
  TestName ->
  -- | Actual
  (forall (s0 :: S). Term s0 a) ->
  -- | Expected
  (forall (s1 :: S). Term s1 a) ->
  TestTree
testEvalEqual :: forall (a :: S -> Type).
TestName
-> (forall (s0 :: S). Term s0 a)
-> (forall (s0 :: S). Term s0 a)
-> TestTree
testEvalEqual TestName
name forall (s0 :: S). Term s0 a
term forall (s0 :: S). Term s0 a
expectedTerm = TestName -> Assertion -> TestTree
testCase TestName
name (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ do
  TestName
actual <- case Config -> (forall (s0 :: S). Term s0 a) -> TermResult
forall (a :: S -> Type).
Config -> (forall (s :: S). Term s a) -> TermResult
evalTermResult Config
NoTracing Term s a
forall (s0 :: S). Term s0 a
term of
    FailedToCompile Text
err -> TestName -> IO TestName
forall a. HasCallStack => TestName -> IO a
assertFailure (TestName -> IO TestName) -> TestName -> IO TestName
forall a b. (a -> b) -> a -> b
$ TestName
"Failed to compile: " TestName -> TestName -> TestName
forall a. Semigroup a => a -> a -> a
<> Text -> TestName
Text.unpack Text
err
    FailedToEvaluate CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
err [Text]
_ -> TestName -> IO TestName
forall a. HasCallStack => TestName -> IO a
assertFailure (TestName -> IO TestName) -> TestName -> IO TestName
forall a b. (a -> b) -> a -> b
$ TestName
"Failed to evaluate: " TestName -> TestName -> TestName
forall a. Semigroup a => a -> a -> a
<> CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
-> TestName
forall a. Show a => a -> TestName
show CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
err
    Evaluated TestName
script [Text]
_ -> TestName -> IO TestName
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure TestName
script
  case Config -> (forall (s0 :: S). Term s0 a) -> TermResult
forall (a :: S -> Type).
Config -> (forall (s :: S). Term s a) -> TermResult
evalTermResult Config
NoTracing Term s a
forall (s0 :: S). Term s0 a
expectedTerm of
    FailedToCompile Text
err -> TestName -> Assertion
forall a. HasCallStack => TestName -> IO a
assertFailure (TestName -> Assertion) -> TestName -> Assertion
forall a b. (a -> b) -> a -> b
$ TestName
"Failed to compile expected term: " TestName -> TestName -> TestName
forall a. Semigroup a => a -> a -> a
<> Text -> TestName
Text.unpack Text
err
    FailedToEvaluate CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
err [Text]
_ -> TestName -> Assertion
forall a. HasCallStack => TestName -> IO a
assertFailure (TestName -> Assertion) -> TestName -> Assertion
forall a b. (a -> b) -> a -> b
$ TestName
"Failed to evaluate expected term: " TestName -> TestName -> TestName
forall a. Semigroup a => a -> a -> a
<> CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
-> TestName
forall a. Show a => a -> TestName
show CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
err
    Evaluated TestName
expected [Text]
_ -> TestName -> TestName -> TestName -> Assertion
forall a.
(Eq a, Show a, HasCallStack) =>
TestName -> a -> a -> Assertion
assertEqual TestName
"" TestName
expected TestName
actual

{- | Assert that term compiled (with specified tracing level and `TracingMode.DetTracing`) and evaluated
without errors produced traces that match expected value. Note that this succeeds even if script
evaluated to error if traces still match

@since 1.0.0
-}
testEvalEqualTraces :: forall (a :: S -> Type). TestName -> (forall (s :: S). Term s a) -> LogLevel -> [Text] -> TestTree
testEvalEqualTraces :: forall (a :: S -> Type).
TestName
-> (forall (s :: S). Term s a) -> LogLevel -> [Text] -> TestTree
testEvalEqualTraces TestName
name forall (s :: S). Term s a
term LogLevel
traceLevel [Text]
expected = TestName -> Assertion -> TestTree
testCase TestName
name (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$
  case Config -> (forall (s :: S). Term s a) -> TermResult
forall (a :: S -> Type).
Config -> (forall (s :: S). Term s a) -> TermResult
evalTermResult (LogLevel -> TracingMode -> Config
Tracing LogLevel
traceLevel TracingMode
DetTracing) Term s a
forall (s :: S). Term s a
term of
    FailedToCompile Text
err -> TestName -> Assertion
forall a. HasCallStack => TestName -> IO a
assertFailure (TestName -> Assertion) -> TestName -> Assertion
forall a b. (a -> b) -> a -> b
$ TestName
"Failed to compile: " TestName -> TestName -> TestName
forall a. Semigroup a => a -> a -> a
<> Text -> TestName
Text.unpack Text
err
    FailedToEvaluate CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
_ [Text]
traces -> TestName -> [Text] -> [Text] -> Assertion
forall a.
(Eq a, Show a, HasCallStack) =>
TestName -> a -> a -> Assertion
assertEqual TestName
"" [Text]
expected [Text]
traces
    Evaluated TestName
_ [Text]
traces -> TestName -> [Text] -> [Text] -> Assertion
forall a.
(Eq a, Show a, HasCallStack) =>
TestName -> a -> a -> Assertion
assertEqual TestName
"" [Text]
expected [Text]
traces

-- | @since 1.0.0
data TermResult
  = FailedToCompile Text
  | FailedToEvaluate (Cek.CekEvaluationException PLC.NamedDeBruijn PLC.DefaultUni PLC.DefaultFun) [Text]
  | Evaluated String [Text]

-- | @since 1.0.0
evalTermResult :: forall (a :: S -> Type). Config -> (forall (s :: S). Term s a) -> TermResult
evalTermResult :: forall (a :: S -> Type).
Config -> (forall (s :: S). Term s a) -> TermResult
evalTermResult Config
config forall (s :: S). Term s a
term =
  case Config -> (forall (s :: S). Term s a) -> Either Text Script
forall (a :: S -> Type).
Config -> (forall (s :: S). Term s a) -> Either Text Script
compile Config
config Term s a
forall (s :: S). Term s a
term of
    Left Text
err -> Text -> TermResult
FailedToCompile Text
err
    Right Script
compiledTerm ->
      case Script
-> (Either
      (CekEvaluationException NamedDeBruijn DefaultUni DefaultFun)
      Script,
    ExBudget, [Text])
evalScriptUnlimited Script
compiledTerm of
        (Left CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
err, ExBudget
_, [Text]
traces) -> CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
-> [Text] -> TermResult
FailedToEvaluate CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
err [Text]
traces
        (Right Script
evaluated, ExBudget
_, [Text]
traces) -> TestName -> [Text] -> TermResult
Evaluated (Script -> TestName
printScript Script
evaluated) [Text]
traces