{-# LANGUAGE DeriveAnyClass    #-}
{-# LANGUAGE LambdaCase        #-}
{-# LANGUAGE OverloadedStrings #-}

module PlutusCore.Evaluation.Machine.CostModelInterface
    ( CostModelParams
    , CekMachineCosts
    , extractCostModelParams
    , applyCostModelParams
    , CostModelApplyError (..)
    )
where

import PlutusCore.Evaluation.Machine.BuiltinCostModel ()
import PlutusCore.Evaluation.Machine.MachineParameters (CostModel (..))
import UntypedPlutusCore.Evaluation.Machine.Cek.CekMachineCosts (CekMachineCosts, cekMachineCostsPrefix)

import Control.Exception
import Control.Monad.Except
import Data.Aeson
import Data.Aeson.Flatten
import Data.HashMap.Strict qualified as HM
import Data.Map qualified as Map
import Data.Map.Merge.Lazy qualified as Map
import Data.Text qualified as Text
import Prettyprinter

{- Note [Cost model parameters]
We want to expose to the ledger some notion of the "cost model
parameters". Intuitively, these should be all the numbers that appear in the
cost model.

However, there are quite a few quirks to deal with.

1. BuiltinCostModel is stuctured.

That is, it's a complex data structure and the numbers in question are often
nested inside it.  To deal with this quickly, we take the ugly approach of
operating on the JSON representation of the model.  We flatten this down into a
simple key-value mapping (see 'flattenObject' and 'unflattenObject'), and then
look only at the numbers.

2. We use CostingIntegers, Aeson uses Data.Scientific.

The numbers in CostModel objects are CostingIntegers, which are usually the
64-bit SatInt type (but Integer on 32-bit machines).  Numerical values in
Aeson-encoded JSON objects are represented as Data.Scientific (Integer mantissa,
Int exponent). We should be able to convert between these types without loss of
precision, except that Scientific numbers of large magnitude will overflow to
SatInt::MaxBound or underflow to SatInt::MinBound.  This is OK because
CostModelParams objects should never contain such large numbers. Any Plutus Core
programs whose cost reaches MaxBound will fail due to excessive resource usage.

3. BuiltinCostModel includes the *type* of the model, which isn't a parameter

We can just strip the type out, but in particular this means that the parameters are
not enough to *construct* a model.  So we punt and say that you can *update* a
model by giving the parameters. So you can take the default model and then
overwrite the parameters, which seems okay.

This is also implemented in a horrible JSON-y way.

4. The implementation is not nice.

Ugly JSON stuff and failure possibilities where there probably shouldn't be any.

5. The overall cost model now includes two components: a model for the internal
costs of the evaluator and a model for built-in evaluation costs.  We just
re-use the technique mentioned above to extract parameters for the evaluator
costs, merging these with the parameters for the builtin cost model to obtain
parameters for the overall model.  To recover cost model components we assume
that every field in the cost model for the evaluator begins with a prefix (eg
"cek") which is does not occur as a prefix of any built-in function, and use
that to split the map of parameters into two maps.

-}

-- See Note [Cost model parameters]
type CostModelParams = Map.Map Text.Text Integer

-- See Note [Cost model parameters]
-- | Extract the model parameters from a model.
extractParams :: ToJSON a => a -> Maybe CostModelParams
extractParams :: a -> Maybe CostModelParams
extractParams a
cm = case a -> Value
forall a. ToJSON a => a -> Value
toJSON a
cm of
    Object Object
o ->
        let
            flattened :: HashMap Text Value
flattened = Object -> HashMap Text Value
objToHm (Object -> HashMap Text Value) -> Object -> HashMap Text Value
forall a b. (a -> b) -> a -> b
$ Text -> Object -> Object
flattenObject Text
"-" Object
o
            usingCostingIntegers :: HashMap Text Integer
usingCostingIntegers = (Value -> Maybe Integer)
-> HashMap Text Value -> HashMap Text Integer
forall v1 v2 k. (v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2
HM.mapMaybe (\case { Number Scientific
n -> Integer -> Maybe Integer
forall a. a -> Maybe a
Just (Integer -> Maybe Integer) -> Integer -> Maybe Integer
forall a b. (a -> b) -> a -> b
$ Scientific -> Integer
forall a b. (RealFrac a, Integral b) => a -> b
ceiling Scientific
n; Value
_ -> Maybe Integer
forall a. Maybe a
Nothing }) HashMap Text Value
flattened
            -- ^ Only (the contents of) the "Just" values are retained in the output map.
            mapified :: CostModelParams
mapified = [(Text, Integer)] -> CostModelParams
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(Text, Integer)] -> CostModelParams)
-> [(Text, Integer)] -> CostModelParams
forall a b. (a -> b) -> a -> b
$ HashMap Text Integer -> [(Text, Integer)]
forall k v. HashMap k v -> [(k, v)]
HM.toList HashMap Text Integer
usingCostingIntegers
        in CostModelParams -> Maybe CostModelParams
forall a. a -> Maybe a
Just CostModelParams
mapified
    Value
_ -> Maybe CostModelParams
forall a. Maybe a
Nothing


-- | The type of errors that 'applyParams' can throw.
data CostModelApplyError =
      CMUnknownParamError Text.Text
      -- ^ a costmodel parameter with the give name does not exist in the costmodel to be applied upon
    | CMInternalReadError
      -- ^ internal error when we are transforming the applyParams' input to json (should not happen)
    | CMInternalWriteError String
      -- ^ internal error when we are transforming the applied params from json with given jsonstring error (should not happen)
    deriving stock Int -> CostModelApplyError -> ShowS
[CostModelApplyError] -> ShowS
CostModelApplyError -> String
(Int -> CostModelApplyError -> ShowS)
-> (CostModelApplyError -> String)
-> ([CostModelApplyError] -> ShowS)
-> Show CostModelApplyError
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [CostModelApplyError] -> ShowS
$cshowList :: [CostModelApplyError] -> ShowS
show :: CostModelApplyError -> String
$cshow :: CostModelApplyError -> String
showsPrec :: Int -> CostModelApplyError -> ShowS
$cshowsPrec :: Int -> CostModelApplyError -> ShowS
Show
    deriving anyclass Show CostModelApplyError
Typeable CostModelApplyError
Typeable CostModelApplyError
-> Show CostModelApplyError
-> (CostModelApplyError -> SomeException)
-> (SomeException -> Maybe CostModelApplyError)
-> (CostModelApplyError -> String)
-> Exception CostModelApplyError
SomeException -> Maybe CostModelApplyError
CostModelApplyError -> String
CostModelApplyError -> SomeException
forall e.
Typeable e
-> Show e
-> (e -> SomeException)
-> (SomeException -> Maybe e)
-> (e -> String)
-> Exception e
displayException :: CostModelApplyError -> String
$cdisplayException :: CostModelApplyError -> String
fromException :: SomeException -> Maybe CostModelApplyError
$cfromException :: SomeException -> Maybe CostModelApplyError
toException :: CostModelApplyError -> SomeException
$ctoException :: CostModelApplyError -> SomeException
$cp2Exception :: Show CostModelApplyError
$cp1Exception :: Typeable CostModelApplyError
Exception

instance Pretty CostModelApplyError where
    pretty :: CostModelApplyError -> Doc ann
pretty = (Doc ann
preamble Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+>) (Doc ann -> Doc ann)
-> (CostModelApplyError -> Doc ann)
-> CostModelApplyError
-> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. \case
        CMUnknownParamError Text
k -> Doc ann
"Unknown cost model parameter:" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Text -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Text
k
        CostModelApplyError
CMInternalReadError      -> Doc ann
"Internal problem occurred upon reading the given cost model parameteres"
        CMInternalWriteError String
str     -> Doc ann
"Internal problem occurred upon generating the applied cost model parameters with JSON error:" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty String
str
      where
          preamble :: Doc ann
preamble = Doc ann
"applyParams error:"

-- See Note [Cost model parameters]
-- | Update a model by overwriting the parameters with the given ones.
applyParams :: (FromJSON a, ToJSON a, MonadError CostModelApplyError m)
            => a
            -> CostModelParams
            -> m a
applyParams :: a -> CostModelParams -> m a
applyParams a
cm CostModelParams
params = case a -> Value
forall a. ToJSON a => a -> Value
toJSON a
cm of
    Object Object
o ->
        let
            usingScientific :: Map Text Value
usingScientific = (Integer -> Value) -> CostModelParams -> Map Text Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Scientific -> Value
Number (Scientific -> Value)
-> (Integer -> Scientific) -> Integer -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Scientific
forall a b. (Integral a, Num b) => a -> b
fromIntegral) CostModelParams
params
            flattened :: Map Text Value
flattened = HashMap Text Value -> Map Text Value
forall a. HashMap Text a -> Map Text a
fromHash (HashMap Text Value -> Map Text Value)
-> HashMap Text Value -> Map Text Value
forall a b. (a -> b) -> a -> b
$ Object -> HashMap Text Value
objToHm (Object -> HashMap Text Value) -> Object -> HashMap Text Value
forall a b. (a -> b) -> a -> b
$ Text -> Object -> Object
flattenObject Text
"-" Object
o
        in do
            -- this is where the overwriting happens
            -- fail when key is in params (left) but not in the model (right)
            Map Text Value
merged <- WhenMissing m Text Value Value
-> WhenMissing m Text Value Value
-> WhenMatched m Text Value Value Value
-> Map Text Value
-> Map Text Value
-> m (Map Text Value)
forall (f :: * -> *) k a c b.
(Applicative f, Ord k) =>
WhenMissing f k a c
-> WhenMissing f k b c
-> WhenMatched f k a b c
-> Map k a
-> Map k b
-> f (Map k c)
Map.mergeA WhenMissing m Text Value Value
forall x y. WhenMissing m Text x y
failMissing WhenMissing m Text Value Value
forall (f :: * -> *) k x. Applicative f => WhenMissing f k x x
Map.preserveMissing ((Text -> Value -> Value -> Value)
-> WhenMatched m Text Value Value Value
forall (f :: * -> *) k x y z.
Applicative f =>
(k -> x -> y -> z) -> WhenMatched f k x y z
Map.zipWithMatched Text -> Value -> Value -> Value
forall p p p. p -> p -> p -> p
leftBiased) Map Text Value
usingScientific Map Text Value
flattened
            let unflattened :: Object
unflattened = Text -> Object -> Object
unflattenObject Text
"-" (Object -> Object) -> Object -> Object
forall a b. (a -> b) -> a -> b
$ HashMap Text Value -> Object
hmToObj (HashMap Text Value -> Object) -> HashMap Text Value -> Object
forall a b. (a -> b) -> a -> b
$ Map Text Value -> HashMap Text Value
forall v. Map Text v -> HashMap Text v
toHash Map Text Value
merged
            case Value -> Result a
forall a. FromJSON a => Value -> Result a
fromJSON (Object -> Value
Object Object
unflattened) of
                Success a
a -> a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a
                Error String
str -> CostModelApplyError -> m a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (CostModelApplyError -> m a) -> CostModelApplyError -> m a
forall a b. (a -> b) -> a -> b
$ String -> CostModelApplyError
CMInternalWriteError String
str
    Value
_ -> CostModelApplyError -> m a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError CostModelApplyError
CMInternalReadError
  where
    toHash :: Map Text v -> HashMap Text v
toHash = [(Text, v)] -> HashMap Text v
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HM.fromList ([(Text, v)] -> HashMap Text v)
-> (Map Text v -> [(Text, v)]) -> Map Text v -> HashMap Text v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Text v -> [(Text, v)]
forall k a. Map k a -> [(k, a)]
Map.toList
    fromHash :: HashMap Text a -> Map Text a
fromHash = [(Text, a)] -> Map Text a
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(Text, a)] -> Map Text a)
-> (HashMap Text a -> [(Text, a)]) -> HashMap Text a -> Map Text a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HashMap Text a -> [(Text, a)]
forall k v. HashMap k v -> [(k, v)]
HM.toList
    -- fail when field missing
    failMissing :: WhenMissing m Text x y
failMissing = (Text -> x -> m y) -> WhenMissing m Text x y
forall (f :: * -> *) k x y.
Applicative f =>
(k -> x -> f y) -> WhenMissing f k x y
Map.traverseMissing ((Text -> x -> m y) -> WhenMissing m Text x y)
-> (Text -> x -> m y) -> WhenMissing m Text x y
forall a b. (a -> b) -> a -> b
$ \ Text
k x
_v -> CostModelApplyError -> m y
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (CostModelApplyError -> m y) -> CostModelApplyError -> m y
forall a b. (a -> b) -> a -> b
$ Text -> CostModelApplyError
CMUnknownParamError Text
k
    -- left-biased merging when key found in both maps
    leftBiased :: p -> p -> p -> p
leftBiased p
_k p
l p
_r = p
l


-- | Parameters for a machine step model and a builtin evaluation model bundled together.
data SplitCostModelParams =
    SplitCostModelParams {
      SplitCostModelParams -> CostModelParams
_machineParams :: CostModelParams
    , SplitCostModelParams -> CostModelParams
_builtinParams :: CostModelParams
    }

-- | Split a CostModelParams object into two subobjects according to some prefix:
-- see item 5 of Note [Cost model parameters].
splitParams :: Text.Text -> CostModelParams -> SplitCostModelParams
splitParams :: Text -> CostModelParams -> SplitCostModelParams
splitParams Text
prefix CostModelParams
params =
    let (CostModelParams
machineparams, CostModelParams
builtinparams) = (Text -> Integer -> Bool)
-> CostModelParams -> (CostModelParams, CostModelParams)
forall k a. (k -> a -> Bool) -> Map k a -> (Map k a, Map k a)
Map.partitionWithKey (\Text
k Integer
_ -> Text -> Text -> Bool
Text.isPrefixOf Text
prefix Text
k) CostModelParams
params
    in CostModelParams -> CostModelParams -> SplitCostModelParams
SplitCostModelParams CostModelParams
machineparams CostModelParams
builtinparams

-- | Given a CostModel, produce a single map containing the parameters from both components
extractCostModelParams
    :: (ToJSON machinecosts, ToJSON builtincosts)
    => CostModel machinecosts builtincosts -> Maybe CostModelParams
extractCostModelParams :: CostModel machinecosts builtincosts -> Maybe CostModelParams
extractCostModelParams CostModel machinecosts builtincosts
model = -- this is using the applicative instance of Maybe
    CostModelParams -> CostModelParams -> CostModelParams
forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union (CostModelParams -> CostModelParams -> CostModelParams)
-> Maybe CostModelParams
-> Maybe (CostModelParams -> CostModelParams)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> machinecosts -> Maybe CostModelParams
forall a. ToJSON a => a -> Maybe CostModelParams
extractParams (CostModel machinecosts builtincosts -> machinecosts
forall machinecosts builtincosts.
CostModel machinecosts builtincosts -> machinecosts
_machineCostModel CostModel machinecosts builtincosts
model) Maybe (CostModelParams -> CostModelParams)
-> Maybe CostModelParams -> Maybe CostModelParams
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> builtincosts -> Maybe CostModelParams
forall a. ToJSON a => a -> Maybe CostModelParams
extractParams (CostModel machinecosts builtincosts -> builtincosts
forall machinecosts builtincosts.
CostModel machinecosts builtincosts -> builtincosts
_builtinCostModel CostModel machinecosts builtincosts
model)

-- | Given a set of cost model parameters, split it into two parts according to
-- some prefix and use those parts to update the components of a cost model.
{- Strictly we don't need to do the splitting: when we call fromJSON in
   applyParams any superfluous objects in the map being decoded will be
   discarded, so we could update both components of the cost model with the
   entire set of parameters without having to worry about splitting the
   parameters on a prefix of the key.  This relies on what appears to be an
   undocumented implementation choice in Aeson though (other JSON decoders (for
   other languages) seem to vary in how unknown fields are handled), so let's be
   explicit. -}
applySplitCostModelParams
    :: (FromJSON evaluatorcosts, FromJSON builtincosts, ToJSON evaluatorcosts, ToJSON builtincosts, MonadError CostModelApplyError m)
    => Text.Text
    -> CostModel evaluatorcosts builtincosts
    -> CostModelParams
    -> m (CostModel evaluatorcosts builtincosts)
applySplitCostModelParams :: Text
-> CostModel evaluatorcosts builtincosts
-> CostModelParams
-> m (CostModel evaluatorcosts builtincosts)
applySplitCostModelParams Text
prefix CostModel evaluatorcosts builtincosts
model CostModelParams
params =
    let SplitCostModelParams CostModelParams
machineparams CostModelParams
builtinparams = Text -> CostModelParams -> SplitCostModelParams
splitParams Text
prefix CostModelParams
params
    in evaluatorcosts
-> builtincosts -> CostModel evaluatorcosts builtincosts
forall machinecosts builtincosts.
machinecosts -> builtincosts -> CostModel machinecosts builtincosts
CostModel (evaluatorcosts
 -> builtincosts -> CostModel evaluatorcosts builtincosts)
-> m evaluatorcosts
-> m (builtincosts -> CostModel evaluatorcosts builtincosts)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> evaluatorcosts -> CostModelParams -> m evaluatorcosts
forall a (m :: * -> *).
(FromJSON a, ToJSON a, MonadError CostModelApplyError m) =>
a -> CostModelParams -> m a
applyParams (CostModel evaluatorcosts builtincosts -> evaluatorcosts
forall machinecosts builtincosts.
CostModel machinecosts builtincosts -> machinecosts
_machineCostModel CostModel evaluatorcosts builtincosts
model) CostModelParams
machineparams
                 m (builtincosts -> CostModel evaluatorcosts builtincosts)
-> m builtincosts -> m (CostModel evaluatorcosts builtincosts)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> builtincosts -> CostModelParams -> m builtincosts
forall a (m :: * -> *).
(FromJSON a, ToJSON a, MonadError CostModelApplyError m) =>
a -> CostModelParams -> m a
applyParams (CostModel evaluatorcosts builtincosts -> builtincosts
forall machinecosts builtincosts.
CostModel machinecosts builtincosts -> builtincosts
_builtinCostModel CostModel evaluatorcosts builtincosts
model) CostModelParams
builtinparams

-- | Update a CostModel for the CEK machine with a given set of parameters,
applyCostModelParams
    :: (FromJSON evaluatorcosts, FromJSON builtincosts, ToJSON evaluatorcosts, ToJSON builtincosts, MonadError CostModelApplyError m)
    => CostModel evaluatorcosts builtincosts
    -> CostModelParams
    -> m (CostModel evaluatorcosts builtincosts)
applyCostModelParams :: CostModel evaluatorcosts builtincosts
-> CostModelParams -> m (CostModel evaluatorcosts builtincosts)
applyCostModelParams = Text
-> CostModel evaluatorcosts builtincosts
-> CostModelParams
-> m (CostModel evaluatorcosts builtincosts)
forall evaluatorcosts builtincosts (m :: * -> *).
(FromJSON evaluatorcosts, FromJSON builtincosts,
 ToJSON evaluatorcosts, ToJSON builtincosts,
 MonadError CostModelApplyError m) =>
Text
-> CostModel evaluatorcosts builtincosts
-> CostModelParams
-> m (CostModel evaluatorcosts builtincosts)
applySplitCostModelParams Text
cekMachineCostsPrefix