{-# LANGUAGE DataKinds                #-}
{-# LANGUAGE GADTs                    #-}
{-# LANGUAGE LambdaCase               #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeFamilies             #-}
{-# LANGUAGE TypeOperators            #-}

{-# LANGUAGE StrictData               #-}

module PlutusCore.Builtin.Runtime where

import PlutusPrelude

import PlutusCore.Builtin.KnownType
import PlutusCore.Evaluation.Machine.ExBudget
import PlutusCore.Evaluation.Machine.ExMemory
import PlutusCore.Evaluation.Machine.Exception

import Control.DeepSeq
import Control.Lens (ix, (^?))
import Control.Monad.Except
import Data.Array
import Data.Kind qualified as GHC (Type)
import NoThunks.Class as NoThunks

-- | Peano numbers. Normally called @Nat@, but that is already reserved by @base@.
data Peano
    = Z
    | S Peano

-- | Same as 'TypeScheme' except this one doesn't contain any evaluation-irrelevant types stuff.
-- @n@ represents the number of term-level arguments that a builtin takes.
data RuntimeScheme n where
    RuntimeSchemeResult :: RuntimeScheme 'Z
    RuntimeSchemeArrow  :: RuntimeScheme n -> RuntimeScheme ('S n)
    RuntimeSchemeAll    :: RuntimeScheme n -> RuntimeScheme n

deriving stock instance Eq   (RuntimeScheme n)
deriving stock instance Show (RuntimeScheme n)

-- we use strictdata, so this is just for the purpose of completeness
instance NFData (RuntimeScheme n) where
    rnf :: RuntimeScheme n -> ()
rnf RuntimeScheme n
r = case RuntimeScheme n
r of
        RuntimeScheme n
RuntimeSchemeResult    -> RuntimeScheme n -> ()
forall a. a -> ()
rwhnf RuntimeScheme n
r
        RuntimeSchemeArrow RuntimeScheme n
arg -> RuntimeScheme n -> ()
forall a. NFData a => a -> ()
rnf RuntimeScheme n
arg
        RuntimeSchemeAll RuntimeScheme n
arg   -> RuntimeScheme n -> ()
forall a. NFData a => a -> ()
rnf RuntimeScheme n
arg

instance NoThunks (RuntimeScheme n) where
    wNoThunks :: Context -> RuntimeScheme n -> IO (Maybe ThunkInfo)
wNoThunks Context
ctx = \case
        RuntimeScheme n
RuntimeSchemeResult  -> Maybe ThunkInfo -> IO (Maybe ThunkInfo)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe ThunkInfo
forall a. Maybe a
Nothing
        RuntimeSchemeArrow RuntimeScheme n
s -> Context -> RuntimeScheme n -> IO (Maybe ThunkInfo)
forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctx RuntimeScheme n
s
        RuntimeSchemeAll RuntimeScheme n
s   -> Context -> RuntimeScheme n -> IO (Maybe ThunkInfo)
forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctx RuntimeScheme n
s
    showTypeOf :: Proxy (RuntimeScheme n) -> String
showTypeOf = String -> Proxy (RuntimeScheme n) -> String
forall a b. a -> b -> a
const String
"PlutusCore.Builtin.Runtime.RuntimeScheme"

-- | Compute the runtime denotation type of a builtin given the type of values and the number of
-- arguments that the builtin takes. A \"runtime denotation type\" is different from a regular
-- denotation type in that a regular one can have any 'ReadKnownIn' type as an argument and can
-- return any 'MakeKnownIn' type, while in a runtime one only @val@ can appear as the type of an
-- argument, as well as in the result type. This is what we get by calling 'readKnown' over each
-- argument of the denotation and calling 'makeKnown' over its result.
type ToRuntimeDenotationType :: GHC.Type -> Peano -> GHC.Type
type family ToRuntimeDenotationType val n where
    ToRuntimeDenotationType val 'Z     = MakeKnownM () val
    -- 'ReadKnownM' is required here only for immediate unlifting, because deferred unlifting
    -- doesn't need the ability to fail in the middle of a builtin application, but having a uniform
    -- interface for both the ways of doing unlifting is way too convenient, hence we decided to pay
    -- the price (about 1-2% of total evaluation time) for now.
    ToRuntimeDenotationType val ('S n) = val -> ReadKnownM () (ToRuntimeDenotationType val n)

-- | Compute the costing type for a builtin given the number of arguments that the builtin takes.
type ToCostingType :: Peano -> GHC.Type
type family ToCostingType n where
    ToCostingType 'Z     = ExBudget
    ToCostingType ('S n) = ExMemory -> ToCostingType n

-- We tried instantiating 'BuiltinMeaning' on the fly and that was slower than precaching
-- 'BuiltinRuntime's.
-- | A 'BuiltinRuntime' represents a possibly partial builtin application.
-- We get an initial 'BuiltinRuntime' representing an empty builtin application (i.e. just the
-- builtin with no arguments) by instantiating (via 'fromBuiltinRuntimeOptions') a
-- 'BuiltinRuntimeOptions'.
--
-- A 'BuiltinRuntime' contains info that is used during evaluation:
--
-- 1. the 'RuntimeScheme' of the uninstantiated part of the builtin. I.e. initially it's the runtime
--      scheme of the whole builtin, but applying or type-instantiating the builtin peels off
--      the corresponding constructor from the runtime scheme
-- 2. the (possibly partially instantiated) runtime denotation
-- 3. the (possibly partially instantiated) costing function
--
-- All the three are in sync in terms of partial instantiatedness due to 'RuntimeScheme' being a
-- GADT and 'ToRuntimeDenotationType' and 'ToCostingType' operating on the index of that GADT.
data BuiltinRuntime val =
    forall n. BuiltinRuntime
        (RuntimeScheme n)
        ~(ToRuntimeDenotationType val n)  -- Must be lazy, because we don't want to compute the
                                          -- denotation when it's fully saturated before figuring
                                          -- out what it's going to cost.
        (ToCostingType n)

instance NoThunks (BuiltinRuntime val) where
    -- Skipping `_denot` in `allNoThunks` check, since it is supposed to be lazy and
    -- is allowed to contain thunks.
    wNoThunks :: Context -> BuiltinRuntime val -> IO (Maybe ThunkInfo)
wNoThunks Context
ctx (BuiltinRuntime RuntimeScheme n
sch ToRuntimeDenotationType val n
_denot ToCostingType n
costing) =
        [IO (Maybe ThunkInfo)] -> IO (Maybe ThunkInfo)
allNoThunks
            [ -- The order here is important: we must first check that `sch` doesn't have
              -- thunks, before inspecting it in `noThunksInCosting`.
              Context -> RuntimeScheme n -> IO (Maybe ThunkInfo)
forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctx RuntimeScheme n
sch
            , Context
-> ToCostingType n -> RuntimeScheme n -> IO (Maybe ThunkInfo)
forall (n :: Peano).
Context
-> ToCostingType n -> RuntimeScheme n -> IO (Maybe ThunkInfo)
noThunksInCosting Context
ctx ToCostingType n
costing RuntimeScheme n
sch
            ]

    showTypeOf :: Proxy (BuiltinRuntime val) -> String
showTypeOf = String -> Proxy (BuiltinRuntime val) -> String
forall a b. a -> b -> a
const String
"PlutusCore.Builtin.Runtime.BuiltinRuntime"

-- | Check whether the given `ToCostingType n` contains thunks. The `RuntimeScheme n` is used to
-- refine `n`.
noThunksInCosting :: NoThunks.Context -> ToCostingType n -> RuntimeScheme n -> IO (Maybe ThunkInfo)
noThunksInCosting :: Context
-> ToCostingType n -> RuntimeScheme n -> IO (Maybe ThunkInfo)
noThunksInCosting Context
ctx ToCostingType n
costing = \case
    -- @n ~ 'Z@, and @ToCostingType n ~ ExBudget@, which should not be a thunk or contain
    -- nested thunks.
    RuntimeScheme n
RuntimeSchemeResult ->
        Context -> ExBudget -> IO (Maybe ThunkInfo)
forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctx ExBudget
ToCostingType n
costing
    RuntimeSchemeArrow RuntimeScheme n
_ ->
        Context -> (ExMemory -> ToCostingType n) -> IO (Maybe ThunkInfo)
forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctx ToCostingType n
ExMemory -> ToCostingType n
costing
    RuntimeSchemeAll RuntimeScheme n
sch ->
        Context
-> ToCostingType n -> RuntimeScheme n -> IO (Maybe ThunkInfo)
forall (n :: Peano).
Context
-> ToCostingType n -> RuntimeScheme n -> IO (Maybe ThunkInfo)
noThunksInCosting Context
ctx ToCostingType n
costing RuntimeScheme n
sch

-- | Determines how to unlift arguments. The difference is that with 'UnliftingImmediate' unlifting
-- is performed immediately after a builtin gets the argument and so can fail immediately too, while
-- with deferred unlifting all arguments are unlifted upon full saturation, hence no failure can
-- occur until that. The former makes it much harder to specify the behaviour of builtins and
-- so 'UnliftingDeferred' is the preferred mode.
data UnliftingMode
    = UnliftingImmediate
    | UnliftingDeferred

-- | A 'BuiltinRuntimeOptions' is a precursor to 'BuiltinRuntime'. One gets the latter from the
-- former by choosing the runtime denotation of the builtin (either '_broImmediateF' for immediate
-- unlifting or '_broDeferredF' for deferred unlifting, see 'UnliftingMode' for details) and by
-- instantiating '_broToExF' with a cost model to get the costing function for the builtin.
data BuiltinRuntimeOptions n val cost = BuiltinRuntimeOptions
    { BuiltinRuntimeOptions n val cost -> RuntimeScheme n
_broRuntimeScheme :: RuntimeScheme n
    , BuiltinRuntimeOptions n val cost -> ToRuntimeDenotationType val n
_broImmediateF    :: ToRuntimeDenotationType val n
    , BuiltinRuntimeOptions n val cost -> ToRuntimeDenotationType val n
_broDeferredF     :: ToRuntimeDenotationType val n
    , BuiltinRuntimeOptions n val cost -> cost -> ToCostingType n
_broToExF         :: cost -> ToCostingType n
    }

-- | Convert a 'BuiltinRuntimeOptions' to a 'BuiltinRuntime' given an 'UnliftingMode' and a cost
-- model.
fromBuiltinRuntimeOptions
    :: UnliftingMode -> cost -> BuiltinRuntimeOptions n val cost -> BuiltinRuntime val
fromBuiltinRuntimeOptions :: UnliftingMode
-> cost -> BuiltinRuntimeOptions n val cost -> BuiltinRuntime val
fromBuiltinRuntimeOptions UnliftingMode
unlMode cost
cost (BuiltinRuntimeOptions RuntimeScheme n
sch ToRuntimeDenotationType val n
immF ToRuntimeDenotationType val n
defF cost -> ToCostingType n
toExF) =
    RuntimeScheme n
-> ToRuntimeDenotationType val n
-> ToCostingType n
-> BuiltinRuntime val
forall val (n :: Peano).
RuntimeScheme n
-> ToRuntimeDenotationType val n
-> ToCostingType n
-> BuiltinRuntime val
BuiltinRuntime RuntimeScheme n
sch ToRuntimeDenotationType val n
f (ToCostingType n -> BuiltinRuntime val)
-> ToCostingType n -> BuiltinRuntime val
forall a b. (a -> b) -> a -> b
$ cost -> ToCostingType n
toExF cost
cost where
        f :: ToRuntimeDenotationType val n
f = case UnliftingMode
unlMode of
                UnliftingMode
UnliftingImmediate -> ToRuntimeDenotationType val n
immF
                UnliftingMode
UnliftingDeferred  -> ToRuntimeDenotationType val n
defF
{-# INLINE fromBuiltinRuntimeOptions #-}

instance NFData (BuiltinRuntime val) where
    rnf :: BuiltinRuntime val -> ()
rnf (BuiltinRuntime RuntimeScheme n
rs ToRuntimeDenotationType val n
f ToCostingType n
exF) = RuntimeScheme n -> ()
forall a. NFData a => a -> ()
rnf RuntimeScheme n
rs () -> () -> ()
`seq` ToRuntimeDenotationType val n
f ToRuntimeDenotationType val n -> () -> ()
`seq` ToCostingType n -> ()
forall a. a -> ()
rwhnf ToCostingType n
exF

-- | A 'BuiltinRuntime' for each builtin from a set of builtins.
newtype BuiltinsRuntime fun val = BuiltinsRuntime
    { BuiltinsRuntime fun val -> Array fun (BuiltinRuntime val)
unBuiltinRuntime :: Array fun (BuiltinRuntime val)
    }

deriving newtype instance (NFData fun) => NFData (BuiltinsRuntime fun val)

instance NoThunks (BuiltinsRuntime fun val) where
    wNoThunks :: Context -> BuiltinsRuntime fun val -> IO (Maybe ThunkInfo)
wNoThunks Context
ctx (BuiltinsRuntime Array fun (BuiltinRuntime val)
arr) = [IO (Maybe ThunkInfo)] -> IO (Maybe ThunkInfo)
allNoThunks (Context -> BuiltinRuntime val -> IO (Maybe ThunkInfo)
forall a. NoThunks a => Context -> a -> IO (Maybe ThunkInfo)
noThunks Context
ctx (BuiltinRuntime val -> IO (Maybe ThunkInfo))
-> [BuiltinRuntime val] -> [IO (Maybe ThunkInfo)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Array fun (BuiltinRuntime val) -> [BuiltinRuntime val]
forall i e. Array i e -> [e]
elems Array fun (BuiltinRuntime val)
arr)
    showTypeOf :: Proxy (BuiltinsRuntime fun val) -> String
showTypeOf = String -> Proxy (BuiltinsRuntime fun val) -> String
forall a b. a -> b -> a
const String
"PlutusCore.Builtin.Runtime.BuiltinsRuntime"

-- | Look up the runtime info of a built-in function during evaluation.
lookupBuiltin
    :: (MonadError (ErrorWithCause err cause) m, AsMachineError err fun, Ix fun)
    => fun -> BuiltinsRuntime fun val -> m (BuiltinRuntime val)
-- @Data.Array@ doesn't seem to have a safe version of @(!)@, hence we use a prism.
lookupBuiltin :: fun -> BuiltinsRuntime fun val -> m (BuiltinRuntime val)
lookupBuiltin fun
fun (BuiltinsRuntime Array fun (BuiltinRuntime val)
env) = case Array fun (BuiltinRuntime val)
env Array fun (BuiltinRuntime val)
-> Getting
     (First (BuiltinRuntime val))
     (Array fun (BuiltinRuntime val))
     (BuiltinRuntime val)
-> Maybe (BuiltinRuntime val)
forall s a. s -> Getting (First a) s a -> Maybe a
^? Index (Array fun (BuiltinRuntime val))
-> Traversal'
     (Array fun (BuiltinRuntime val))
     (IxValue (Array fun (BuiltinRuntime val)))
forall m. Ixed m => Index m -> Traversal' m (IxValue m)
ix fun
Index (Array fun (BuiltinRuntime val))
fun of
    Maybe (BuiltinRuntime val)
Nothing  -> AReview err (MachineError fun)
-> MachineError fun -> Maybe cause -> m (BuiltinRuntime val)
forall exc e t term (m :: * -> *) x.
(exc ~ ErrorWithCause e term, MonadError exc m) =>
AReview e t -> t -> Maybe term -> m x
throwingWithCause AReview err (MachineError fun)
forall r fun. AsMachineError r fun => Prism' r (MachineError fun)
_MachineError (fun -> MachineError fun
forall fun. fun -> MachineError fun
UnknownBuiltin fun
fun) Maybe cause
forall a. Maybe a
Nothing
    Just BuiltinRuntime val
bri -> BuiltinRuntime val -> m (BuiltinRuntime val)
forall (f :: * -> *) a. Applicative f => a -> f a
pure BuiltinRuntime val
bri