{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies  #-}
{-# LANGUAGE TypeOperators #-}
module PlutusLedgerApi.Test.Common.EvaluationContext
    ( MCostModel
    , MCekMachineCosts
    , MBuiltinCostModel
    , toMCostModel
    , extractCostModelParamsLedgerOrder
    ) where

import PlutusCore.Evaluation.Machine.BuiltinCostModel
import PlutusCore.Evaluation.Machine.CostModelInterface
import PlutusCore.Evaluation.Machine.MachineParameters
import PlutusPrelude
import UntypedPlutusCore.Evaluation.Machine.Cek.CekMachineCosts

import PlutusLedgerApi.Common as Common

import Barbies
import Data.Functor.Identity
import Data.Int (Int64)
import Data.Map qualified as Map

-- A lifted cost model to `Maybe`, so we can easily clear some of its fields when extracting JSON.
type MCostModel = CostModel MCekMachineCosts MBuiltinCostModel

type MCekMachineCosts = CekMachineCostsBase Maybe

type MBuiltinCostModel = BuiltinCostModelBase MCostingFun

-- | A helper function to lift to a "full" `MCostModel`, by mapping *all* of its fields to `Just`.
-- The fields can be later on cleared, by assigning them to `Nothing`.
toMCostModel :: CostModel CekMachineCosts BuiltinCostModel
             -> MCostModel
toMCostModel :: CostModel CekMachineCosts BuiltinCostModel -> MCostModel
toMCostModel CostModel CekMachineCosts BuiltinCostModel
cm =
   CostModel CekMachineCosts BuiltinCostModel
cm
   CostModel CekMachineCosts BuiltinCostModel
-> (CostModel CekMachineCosts BuiltinCostModel
    -> CostModel (CekMachineCostsBase Maybe) BuiltinCostModel)
-> CostModel (CekMachineCostsBase Maybe) BuiltinCostModel
forall a b. a -> (a -> b) -> b
& (CekMachineCosts -> Identity (CekMachineCostsBase Maybe))
-> CostModel CekMachineCosts BuiltinCostModel
-> Identity
     (CostModel (CekMachineCostsBase Maybe) BuiltinCostModel)
forall machinecosts1 builtincosts machinecosts2 (f :: * -> *).
Functor f =>
(machinecosts1 -> f machinecosts2)
-> CostModel machinecosts1 builtincosts
-> f (CostModel machinecosts2 builtincosts)
machineCostModel
   ((CekMachineCosts -> Identity (CekMachineCostsBase Maybe))
 -> CostModel CekMachineCosts BuiltinCostModel
 -> Identity
      (CostModel (CekMachineCostsBase Maybe) BuiltinCostModel))
-> (CekMachineCosts -> CekMachineCostsBase Maybe)
-> CostModel CekMachineCosts BuiltinCostModel
-> CostModel (CekMachineCostsBase Maybe) BuiltinCostModel
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (forall a. Identity a -> Maybe a)
-> CekMachineCosts -> CekMachineCostsBase Maybe
forall k (b :: (k -> *) -> *) (f :: k -> *) (g :: k -> *).
FunctorB b =>
(forall (a :: k). f a -> g a) -> b f -> b g
forall (f :: * -> *) (g :: * -> *).
(forall a. f a -> g a)
-> CekMachineCostsBase f -> CekMachineCostsBase g
bmap (a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> (Identity a -> a) -> Identity a -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Identity a -> a
forall a. Identity a -> a
runIdentity)
   CostModel (CekMachineCostsBase Maybe) BuiltinCostModel
-> (CostModel (CekMachineCostsBase Maybe) BuiltinCostModel
    -> MCostModel)
-> MCostModel
forall a b. a -> (a -> b) -> b
& (BuiltinCostModel -> Identity MBuiltinCostModel)
-> CostModel (CekMachineCostsBase Maybe) BuiltinCostModel
-> Identity MCostModel
forall machinecosts builtincosts1 builtincosts2 (f :: * -> *).
Functor f =>
(builtincosts1 -> f builtincosts2)
-> CostModel machinecosts builtincosts1
-> f (CostModel machinecosts builtincosts2)
builtinCostModel
   ((BuiltinCostModel -> Identity MBuiltinCostModel)
 -> CostModel (CekMachineCostsBase Maybe) BuiltinCostModel
 -> Identity MCostModel)
-> (BuiltinCostModel -> MBuiltinCostModel)
-> CostModel (CekMachineCostsBase Maybe) BuiltinCostModel
-> MCostModel
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (forall a. CostingFun a -> MCostingFun a)
-> BuiltinCostModel -> MBuiltinCostModel
forall k (b :: (k -> *) -> *) (f :: k -> *) (g :: k -> *).
FunctorB b =>
(forall (a :: k). f a -> g a) -> b f -> b g
forall (f :: * -> *) (g :: * -> *).
(forall a. f a -> g a)
-> BuiltinCostModelBase f -> BuiltinCostModelBase g
bmap (Maybe (CostingFun a) -> MCostingFun a
forall a. Maybe (CostingFun a) -> MCostingFun a
MCostingFun (Maybe (CostingFun a) -> MCostingFun a)
-> (CostingFun a -> Maybe (CostingFun a))
-> CostingFun a
-> MCostingFun a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CostingFun a -> Maybe (CostingFun a)
forall a. a -> Maybe a
Just)

{- | A variant of `extractCostModelParams` to make a mapping of params not in alphabetical order,
but in the `ParamName` order, i.e. the order expected by the ledger.

Here, overconstrained to `MCostModel`, but it could also work with `CostModel mcosts bcosts`.
-}
extractCostModelParamsLedgerOrder :: (Common.IsParamName p, Ord p)
                                  => MCostModel
                                  -> Maybe (Map.Map p Int64)
extractCostModelParamsLedgerOrder :: forall p.
(IsParamName p, Ord p) =>
MCostModel -> Maybe (Map p Int64)
extractCostModelParamsLedgerOrder =
    MCostModel -> Maybe CostModelParams
extractInAlphaOrder
     (MCostModel -> Maybe CostModelParams)
-> (CostModelParams -> Maybe (Map p Int64))
-> MCostModel
-> Maybe (Map p Int64)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> CostModelParams -> Maybe (Map p Int64)
forall {a}. Map Text a -> Maybe (Map p a)
toLedgerOrder
    where
      extractInAlphaOrder :: MCostModel -> Maybe CostModelParams
extractInAlphaOrder = MCostModel -> Maybe CostModelParams
forall machinecosts builtincosts.
(ToJSON machinecosts, ToJSON builtincosts) =>
CostModel machinecosts builtincosts -> Maybe CostModelParams
extractCostModelParams
      toLedgerOrder :: Map Text a -> Maybe (Map p a)
toLedgerOrder = (Text -> Maybe p) -> Map Text a -> Maybe (Map p a)
forall (m :: * -> *) k2 k1 a.
(Monad m, Ord k2) =>
(k1 -> m k2) -> Map k1 a -> m (Map k2 a)
mapKeysM Text -> Maybe p
forall a. IsParamName a => Text -> Maybe a
readParamName

      mapKeysM :: (Monad m, Ord k2) => (k1 -> m k2) -> Map.Map k1 a -> m (Map.Map k2 a)
      mapKeysM :: forall (m :: * -> *) k2 k1 a.
(Monad m, Ord k2) =>
(k1 -> m k2) -> Map k1 a -> m (Map k2 a)
mapKeysM = ([(k1, a)] -> m [(k2, a)]) -> Map k1 a -> m (Map k2 a)
forall {f :: * -> *} {k} {k} {a} {a}.
(Functor f, Ord k) =>
([(k, a)] -> f [(k, a)]) -> Map k a -> f (Map k a)
viaListM (([(k1, a)] -> m [(k2, a)]) -> Map k1 a -> m (Map k2 a))
-> ((k1 -> m k2) -> [(k1, a)] -> m [(k2, a)])
-> (k1 -> m k2)
-> Map k1 a
-> m (Map k2 a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k1, a) -> m (k2, a)) -> [(k1, a)] -> m [(k2, a)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (((k1, a) -> m (k2, a)) -> [(k1, a)] -> m [(k2, a)])
-> ((k1 -> m k2) -> (k1, a) -> m (k2, a))
-> (k1 -> m k2)
-> [(k1, a)]
-> m [(k2, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (k1 -> m k2) -> (k1, a) -> m (k2, a)
forall {f :: * -> *} {t} {a} {t}.
Functor f =>
(t -> f a) -> (t, t) -> f (a, t)
firstM

      viaListM :: ([(k, a)] -> f [(k, a)]) -> Map k a -> f (Map k a)
viaListM [(k, a)] -> f [(k, a)]
op = ([(k, a)] -> Map k a) -> f [(k, a)] -> f (Map k a)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(k, a)] -> Map k a
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList (f [(k, a)] -> f (Map k a))
-> (Map k a -> f [(k, a)]) -> Map k a -> f (Map k a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(k, a)] -> f [(k, a)]
op ([(k, a)] -> f [(k, a)])
-> (Map k a -> [(k, a)]) -> Map k a -> f [(k, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k a -> [(k, a)]
forall k a. Map k a -> [(k, a)]
Map.toList
      firstM :: (t -> f a) -> (t, t) -> f (a, t)
firstM t -> f a
f (t
k,t
v) = (,t
v) (a -> (a, t)) -> f a -> f (a, t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> t -> f a
f t
k