{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Transform.Inline.Spec where

import Control.Monad.Reader (runReaderT)
import Control.Monad.State (runStateT)
import PlutusCore.Annotation (Inline (MayInline))
import PlutusCore.Default (DefaultFun (..), DefaultUni)
import PlutusCore.Name.Unique (Name)
import PlutusCore.Quote (runQuote)
import PlutusPrelude (def)
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (Assertion, assertBool, testCase)
import Transform.Lib (T, app, builtin, case_, delay, lam, name, var)
import UntypedPlutusCore.AstSize (AstSize (..))
import UntypedPlutusCore.Core (Term (..))
import UntypedPlutusCore.Transform.Inline
  ( Ann
  , InlineHints (..)
  , InlineInfo (..)
  , InlineM
  , S (..)
  , Subst (..)
  , TermEnv (..)
  , effectSafe
  , isFirstVarBeforeEffects
  , isStrictIn
  )

test_inline :: TestTree
test_inline :: TestTree
test_inline =
  TestName -> [TestTree] -> TestTree
testGroup
    TestName
"Inline"
    [ TestName -> Assertion -> TestTree
testCase
        TestName
"var is before or after effects"
        Assertion
testVarBeforeAfterEffects
    , TestName -> [TestTree] -> TestTree
testGroup
        TestName
"isStrictIn"
        [ TestName -> Assertion -> TestTree
testCase
            TestName
"a var is delayed if it's inside a delay"
            Assertion
testVarIsEventuallyEvaluatedDelay
        , TestName -> Assertion -> TestTree
testCase
            TestName
"a var is delayed if it's inside a lambda"
            Assertion
testVarIsEventuallyEvaluatedLambda
        , TestName -> Assertion -> TestTree
testCase
            TestName
"a var is delayed if it's inside a case branch"
            Assertion
testVarIsEventuallyEvaluatedCaseBranch
        ]
    , TestName -> [TestTree] -> TestTree
testGroup
        TestName
"effectSafe"
        [ TestName -> Assertion -> TestTree
testCase TestName
"with preserved logs" Assertion
testEffectSafePreservedLogs
        , TestName -> Assertion -> TestTree
testCase TestName
"without preserved logs" Assertion
testEffectSafeWithoutPreservedLogs
        ]
    ]

testVarBeforeAfterEffects :: Assertion
testVarBeforeAfterEffects :: Assertion
testVarBeforeAfterEffects = do
  HasCallStack => TestName -> Bool -> Assertion
TestName -> Bool -> Assertion
assertBool TestName
"a is evaluated before effects" do
    BuiltinSemanticsVariant DefaultFun
-> Name -> Term Name DefaultUni DefaultFun () -> Bool
forall name (uni :: * -> *) fun ann.
InliningConstraints name uni fun =>
BuiltinSemanticsVariant fun
-> name -> Term name uni fun ann -> Bool
isFirstVarBeforeEffects BuiltinSemanticsVariant DefaultFun
forall a. Default a => a
def (TestName -> Name
name TestName
"a") Term Name DefaultUni DefaultFun ()
term
  HasCallStack => TestName -> Bool -> Assertion
TestName -> Bool -> Assertion
assertBool TestName
"b is evaluated before effects" do
    BuiltinSemanticsVariant DefaultFun
-> Name -> Term Name DefaultUni DefaultFun () -> Bool
forall name (uni :: * -> *) fun ann.
InliningConstraints name uni fun =>
BuiltinSemanticsVariant fun
-> name -> Term name uni fun ann -> Bool
isFirstVarBeforeEffects BuiltinSemanticsVariant DefaultFun
forall a. Default a => a
def (TestName -> Name
name TestName
"b") Term Name DefaultUni DefaultFun ()
term
  HasCallStack => TestName -> Bool -> Assertion
TestName -> Bool -> Assertion
assertBool TestName
"c is not evaluated before effects" (Bool -> Assertion) -> Bool -> Assertion
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not do
    BuiltinSemanticsVariant DefaultFun
-> Name -> Term Name DefaultUni DefaultFun () -> Bool
forall name (uni :: * -> *) fun ann.
InliningConstraints name uni fun =>
BuiltinSemanticsVariant fun
-> name -> Term name uni fun ann -> Bool
isFirstVarBeforeEffects BuiltinSemanticsVariant DefaultFun
forall a. Default a => a
def (TestName -> Name
name TestName
"c") Term Name DefaultUni DefaultFun ()
term
  where
    term :: Term Name DefaultUni DefaultFun ()
    term :: Term Name DefaultUni DefaultFun ()
term =
      {- Evaluation order:

        1. pure work-free: a
        2. pure work-free: b
        3. impure? maybe work?: plus a b
        4. pure work-free: c
        5. impure? maybe work?: plus (plus a b) c
      -}
      Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
plus (Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
plus (TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"a") (TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"b")) (TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"c")

testVarIsEventuallyEvaluatedDelay :: Assertion
testVarIsEventuallyEvaluatedDelay :: Assertion
testVarIsEventuallyEvaluatedDelay = do
  HasCallStack => TestName -> Bool -> Assertion
TestName -> Bool -> Assertion
assertBool TestName
"var 'a' is not eventually evaluated in delay" (Bool -> Assertion) -> Bool -> Assertion
forall a b. (a -> b) -> a -> b
$
    Bool -> Bool
not (Name -> Term Name DefaultUni DefaultFun () -> Bool
forall name (uni :: * -> *) fun a.
Eq name =>
name -> Term name uni fun a -> Bool
isStrictIn (TestName -> Name
name TestName
"a") Term Name DefaultUni DefaultFun ()
term)
  HasCallStack => TestName -> Bool -> Assertion
TestName -> Bool -> Assertion
assertBool TestName
"var 'b' is eventually evaluated outside of the delay" (Bool -> Assertion) -> Bool -> Assertion
forall a b. (a -> b) -> a -> b
$
    Name -> Term Name DefaultUni DefaultFun () -> Bool
forall name (uni :: * -> *) fun a.
Eq name =>
name -> Term name uni fun a -> Bool
isStrictIn (TestName -> Name
name TestName
"b") Term Name DefaultUni DefaultFun ()
term
  HasCallStack => TestName -> Bool -> Assertion
TestName -> Bool -> Assertion
assertBool TestName
"it's not known if var 'c' is eventually evaluated" (Bool -> Assertion) -> Bool -> Assertion
forall a b. (a -> b) -> a -> b
$
    Bool -> Bool
not (Name -> Term Name DefaultUni DefaultFun () -> Bool
forall name (uni :: * -> *) fun a.
Eq name =>
name -> Term name uni fun a -> Bool
isStrictIn (TestName -> Name
name TestName
"c") Term Name DefaultUni DefaultFun ()
term)
  where
    term :: Term Name DefaultUni DefaultFun ()
    term :: Term Name DefaultUni DefaultFun ()
term = Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
delay (TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"a" Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
`plus` TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"b") Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
`plus` TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"b"

testVarIsEventuallyEvaluatedLambda :: Assertion
testVarIsEventuallyEvaluatedLambda :: Assertion
testVarIsEventuallyEvaluatedLambda = do
  HasCallStack => TestName -> Bool -> Assertion
TestName -> Bool -> Assertion
assertBool TestName
"var 'a' is not eventually evaluated in lambda body" (Bool -> Assertion) -> Bool -> Assertion
forall a b. (a -> b) -> a -> b
$
    Bool -> Bool
not (Name -> Term Name DefaultUni DefaultFun () -> Bool
forall name (uni :: * -> *) fun a.
Eq name =>
name -> Term name uni fun a -> Bool
isStrictIn (TestName -> Name
name TestName
"a") Term Name DefaultUni DefaultFun ()
term)
  HasCallStack => TestName -> Bool -> Assertion
TestName -> Bool -> Assertion
assertBool TestName
"var 'c' is eventually evaluated outside of the lambda" (Bool -> Assertion) -> Bool -> Assertion
forall a b. (a -> b) -> a -> b
$
    Name -> Term Name DefaultUni DefaultFun () -> Bool
forall name (uni :: * -> *) fun a.
Eq name =>
name -> Term name uni fun a -> Bool
isStrictIn (TestName -> Name
name TestName
"c") Term Name DefaultUni DefaultFun ()
term
  HasCallStack => TestName -> Bool -> Assertion
TestName -> Bool -> Assertion
assertBool TestName
"it's not known if var 'd' is eventually evaluated" (Bool -> Assertion) -> Bool -> Assertion
forall a b. (a -> b) -> a -> b
$
    Bool -> Bool
not (Name -> Term Name DefaultUni DefaultFun () -> Bool
forall name (uni :: * -> *) fun a.
Eq name =>
name -> Term name uni fun a -> Bool
isStrictIn (TestName -> Name
name TestName
"d") Term Name DefaultUni DefaultFun ()
term)
  where
    term :: Term Name DefaultUni DefaultFun ()
    term :: Term Name DefaultUni DefaultFun ()
term = TestName
-> Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
lam TestName
"b" (TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"a" Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
`plus` TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"c") Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
`app` TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"c"

testVarIsEventuallyEvaluatedCaseBranch :: Assertion
testVarIsEventuallyEvaluatedCaseBranch :: Assertion
testVarIsEventuallyEvaluatedCaseBranch = do
  HasCallStack => TestName -> Bool -> Assertion
TestName -> Bool -> Assertion
assertBool TestName
"var 'a' is not eventually evaluated in case branch" (Bool -> Assertion) -> Bool -> Assertion
forall a b. (a -> b) -> a -> b
$
    Bool -> Bool
not (Name -> Term Name DefaultUni DefaultFun () -> Bool
forall name (uni :: * -> *) fun a.
Eq name =>
name -> Term name uni fun a -> Bool
isStrictIn (TestName -> Name
name TestName
"a") Term Name DefaultUni DefaultFun ()
term)
  HasCallStack => TestName -> Bool -> Assertion
TestName -> Bool -> Assertion
assertBool TestName
"var 'b' is eventually evaluated outside of the case branch" (Bool -> Assertion) -> Bool -> Assertion
forall a b. (a -> b) -> a -> b
$
    Name -> Term Name DefaultUni DefaultFun () -> Bool
forall name (uni :: * -> *) fun a.
Eq name =>
name -> Term name uni fun a -> Bool
isStrictIn (TestName -> Name
name TestName
"b") Term Name DefaultUni DefaultFun ()
term
  HasCallStack => TestName -> Bool -> Assertion
TestName -> Bool -> Assertion
assertBool TestName
"it is not known if var 'd' is eventually evaluated" (Bool -> Assertion) -> Bool -> Assertion
forall a b. (a -> b) -> a -> b
$
    Bool -> Bool
not (Name -> Term Name DefaultUni DefaultFun () -> Bool
forall name (uni :: * -> *) fun a.
Eq name =>
name -> Term name uni fun a -> Bool
isStrictIn (TestName -> Name
name TestName
"d") Term Name DefaultUni DefaultFun ()
term)
  where
    term :: Term Name DefaultUni DefaultFun ()
    term :: Term Name DefaultUni DefaultFun ()
term = Term Name DefaultUni DefaultFun ()
-> [Term Name DefaultUni DefaultFun ()]
-> Term Name DefaultUni DefaultFun ()
case_ (TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"b") [TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"a", TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"b", TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"c"]

testEffectSafePreservedLogs :: Assertion
testEffectSafePreservedLogs :: Assertion
testEffectSafePreservedLogs = do
  HasCallStack => TestName -> Bool -> Assertion
TestName -> Bool -> Assertion
assertBool TestName
"an immediately eval'd var is not \"effect safe\"" (Bool -> Assertion) -> Bool -> Assertion
forall a b. (a -> b) -> a -> b
$
    InlineM Name DefaultUni DefaultFun () Bool -> Bool
forall r. InlineM Name DefaultUni DefaultFun () r -> r
runInlineWithLogging (Bool -> Bool
not (Bool -> Bool)
-> InlineM Name DefaultUni DefaultFun () Bool
-> InlineM Name DefaultUni DefaultFun () Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Term Name DefaultUni DefaultFun ()
-> Name -> Bool -> InlineM Name DefaultUni DefaultFun () Bool
forall name (uni :: * -> *) fun a b.
InliningConstraints name uni fun =>
Term name uni fun a -> name -> Bool -> InlineM name uni fun b Bool
effectSafe Term Name DefaultUni DefaultFun ()
term (TestName -> Name
name TestName
"c") Bool
False)
  HasCallStack => TestName -> Bool -> Assertion
TestName -> Bool -> Assertion
assertBool TestName
"a var before effects is \"effect safe\"" (Bool -> Assertion) -> Bool -> Assertion
forall a b. (a -> b) -> a -> b
$
    InlineM Name DefaultUni DefaultFun () Bool -> Bool
forall r. InlineM Name DefaultUni DefaultFun () r -> r
runInlineWithLogging (Term Name DefaultUni DefaultFun ()
-> Name -> Bool -> InlineM Name DefaultUni DefaultFun () Bool
forall name (uni :: * -> *) fun a b.
InliningConstraints name uni fun =>
Term name uni fun a -> name -> Bool -> InlineM name uni fun b Bool
effectSafe Term Name DefaultUni DefaultFun ()
term (TestName -> Name
name TestName
"a") Bool
False)
  where
    term :: Term Name DefaultUni DefaultFun ()
    term :: Term Name DefaultUni DefaultFun ()
term = (TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"a" Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
`plus` TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"b") Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
`plus` TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"c"

testEffectSafeWithoutPreservedLogs :: Assertion
testEffectSafeWithoutPreservedLogs :: Assertion
testEffectSafeWithoutPreservedLogs = do
  HasCallStack => TestName -> Bool -> Assertion
TestName -> Bool -> Assertion
assertBool TestName
"an immediately eval'd var is \"effect safe\"" (Bool -> Assertion) -> Bool -> Assertion
forall a b. (a -> b) -> a -> b
$
    InlineM Name DefaultUni DefaultFun () Bool -> Bool
forall r. InlineM Name DefaultUni DefaultFun () r -> r
runInlineWithoutLogging (Term Name DefaultUni DefaultFun ()
-> Name -> Bool -> InlineM Name DefaultUni DefaultFun () Bool
forall name (uni :: * -> *) fun a b.
InliningConstraints name uni fun =>
Term name uni fun a -> name -> Bool -> InlineM name uni fun b Bool
effectSafe Term Name DefaultUni DefaultFun ()
term (TestName -> Name
name TestName
"c") Bool
False)
  HasCallStack => TestName -> Bool -> Assertion
TestName -> Bool -> Assertion
assertBool TestName
"a var before effects is \"effect safe\"" (Bool -> Assertion) -> Bool -> Assertion
forall a b. (a -> b) -> a -> b
$
    InlineM Name DefaultUni DefaultFun () Bool -> Bool
forall r. InlineM Name DefaultUni DefaultFun () r -> r
runInlineWithoutLogging (Term Name DefaultUni DefaultFun ()
-> Name -> Bool -> InlineM Name DefaultUni DefaultFun () Bool
forall name (uni :: * -> *) fun a b.
InliningConstraints name uni fun =>
Term name uni fun a -> name -> Bool -> InlineM name uni fun b Bool
effectSafe Term Name DefaultUni DefaultFun ()
term (TestName -> Name
name TestName
"a") Bool
False)
  where
    term :: Term Name DefaultUni DefaultFun ()
    term :: Term Name DefaultUni DefaultFun ()
term = (TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"a" Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
`plus` TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"b") Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
`plus` TestName -> Term Name DefaultUni DefaultFun ()
var TestName
"c"

--------------------------------------------------------------------------------
-- InlineM runner --------------------------------------------------------------

runInlineWithoutLogging :: InlineM Name DefaultUni DefaultFun () r -> r
runInlineWithoutLogging :: forall r. InlineM Name DefaultUni DefaultFun () r -> r
runInlineWithoutLogging = Bool -> InlineM Name DefaultUni DefaultFun () r -> r
forall r. Bool -> InlineM Name DefaultUni DefaultFun () r -> r
runInlineM Bool
False

runInlineWithLogging :: InlineM Name DefaultUni DefaultFun () r -> r
runInlineWithLogging :: forall r. InlineM Name DefaultUni DefaultFun () r -> r
runInlineWithLogging = Bool -> InlineM Name DefaultUni DefaultFun () r -> r
forall r. Bool -> InlineM Name DefaultUni DefaultFun () r -> r
runInlineM Bool
True

runInlineM :: Bool -> InlineM Name DefaultUni DefaultFun () r -> r
runInlineM :: forall r. Bool -> InlineM Name DefaultUni DefaultFun () r -> r
runInlineM Bool
preserveLogging InlineM Name DefaultUni DefaultFun () r
m = r
result
  where
    (r
result, S Name DefaultUni DefaultFun (Ann ())
_finalState) =
      Quote (r, S Name DefaultUni DefaultFun (Ann ()))
-> (r, S Name DefaultUni DefaultFun (Ann ()))
forall a. Quote a -> a
runQuote (StateT (S Name DefaultUni DefaultFun (Ann ())) Quote r
-> S Name DefaultUni DefaultFun (Ann ())
-> Quote (r, S Name DefaultUni DefaultFun (Ann ()))
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (InlineM Name DefaultUni DefaultFun () r
-> InlineInfo Name DefaultFun ()
-> StateT (S Name DefaultUni DefaultFun (Ann ())) Quote r
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT InlineM Name DefaultUni DefaultFun () r
m InlineInfo Name DefaultFun ()
inlineInfo) S Name DefaultUni DefaultFun (Ann ())
initialState)
    inlineInfo :: InlineInfo Name DefaultFun ()
    inlineInfo :: InlineInfo Name DefaultFun ()
inlineInfo =
      InlineInfo
        { _iiUsages :: Usages
_iiUsages = Usages
forall a. Monoid a => a
mempty
        , _iiHints :: InlineHints Name ()
_iiHints = (() -> Name -> Inline) -> InlineHints Name ()
forall name a. (a -> name -> Inline) -> InlineHints name a
InlineHints \()
_ann Name
_name -> Inline
MayInline
        , _iiBuiltinSemanticsVariant :: BuiltinSemanticsVariant DefaultFun
_iiBuiltinSemanticsVariant = BuiltinSemanticsVariant DefaultFun
forall a. Default a => a
def
        , _iiInlineConstants :: Bool
_iiInlineConstants = Bool
True
        , _iiInlineUnconditionalGrowth :: AstSize
_iiInlineUnconditionalGrowth = Integer -> AstSize
AstSize Integer
3
        , _iiInlineCallsiteGrowth :: AstSize
_iiInlineCallsiteGrowth = Integer -> AstSize
AstSize Integer
1_000_000
        , _iiPreserveLogging :: Bool
_iiPreserveLogging = Bool
preserveLogging
        }
    initialState :: S Name DefaultUni DefaultFun (Ann ())
    initialState :: S Name DefaultUni DefaultFun (Ann ())
initialState = S {_subst :: Subst Name DefaultUni DefaultFun (Ann ())
_subst = TermEnv Name DefaultUni DefaultFun (Ann ())
-> Subst Name DefaultUni DefaultFun (Ann ())
forall name (uni :: * -> *) fun a.
TermEnv name uni fun a -> Subst name uni fun a
Subst (UniqueMap
  TermUnique (InlineTerm Name DefaultUni DefaultFun (Ann ()))
-> TermEnv Name DefaultUni DefaultFun (Ann ())
forall name (uni :: * -> *) fun a.
UniqueMap TermUnique (InlineTerm name uni fun a)
-> TermEnv name uni fun a
TermEnv UniqueMap
  TermUnique (InlineTerm Name DefaultUni DefaultFun (Ann ()))
forall a. Monoid a => a
mempty), _vars :: UniqueMap TermUnique (VarInfo Name DefaultUni DefaultFun (Ann ()))
_vars = UniqueMap TermUnique (VarInfo Name DefaultUni DefaultFun (Ann ()))
forall a. Monoid a => a
mempty}

--------------------------------------------------------------------------------
-- Local helpers ---------------------------------------------------------------

plus :: T -> T -> T
plus :: Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
plus Term Name DefaultUni DefaultFun ()
i Term Name DefaultUni DefaultFun ()
j = DefaultFun -> Term Name DefaultUni DefaultFun ()
builtin DefaultFun
AddInteger Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
`app` Term Name DefaultUni DefaultFun ()
i Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
`app` Term Name DefaultUni DefaultFun ()
j