Add code review suggestions

* Override cost models when injecting into test state (if extra config
cost models are provided)
* Preserve `AlonzoGenesis` JSON and binary representations but change
the Haskell side representation

Author: Lucsanszky

eras/alonzo/impl/src/Cardano/Ledger/Alonzo/Genesis.hs

@@ -4,14 +4,12 @@
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE LambdaCase #-}
-{-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE UndecidableInstances #-}
-{-# LANGUAGE ViewPatterns #-}
 {-# OPTIONS_GHC -Wno-orphans #-}
 
 module Cardano.Ledger.Alonzo.Genesis (
@@ -43,9 +41,12 @@ import Cardano.Ledger.Alonzo.Scripts (
   CostModels,
   ExUnits (..),
   Prices (..),
+  costModelsValid,
   decodeCostModel,
   decodeCostModelsLenient,
   encodeCostModel,
+  flattenCostModels,
+  mkCostModels,
  )
 import Cardano.Ledger.BaseTypes (KeyValuePairs (..), ToKeyValuePairs (..))
 import Cardano.Ledger.Binary (
@@ -59,20 +60,22 @@ import Cardano.Ledger.Binary (
 import Cardano.Ledger.Binary.Coders
 import Cardano.Ledger.Core
 import Cardano.Ledger.Genesis (EraGenesis (..))
-import Cardano.Ledger.Plutus.CostModels (parseCostModelAsArray, parseCostModels)
-import Cardano.Ledger.Plutus.Language (Language (..))
+import Cardano.Ledger.Plutus (Language (PlutusV1))
+import Cardano.Ledger.Plutus.CostModels (parseCostModels)
 import Control.DeepSeq (NFData)
 import Data.Aeson (FromJSON (..), ToJSON (..), (.:), (.:?), (.=))
 import qualified Data.Aeson as Aeson
 import Data.Functor.Identity (Identity)
+import qualified Data.List as List
+import qualified Data.Map.Strict as Map
 import GHC.Generics (Generic)
 import NoThunks.Class (NoThunks)
 import Numeric.Natural (Natural)
 
 -- | All configuration that is necessary to bootstrap AlonzoEra from ShelleyGenesis
 data AlonzoGenesis = AlonzoGenesisWrapper
   { unAlonzoGenesisWrapper :: UpgradeAlonzoPParams Identity
-  , extraConfig :: AlonzoExtraConfig
+  , extraConfig :: Maybe AlonzoExtraConfig
   }
   deriving stock (Eq, Show, Generic)
   deriving (ToJSON) via KeyValuePairs AlonzoGenesis
@@ -123,7 +126,7 @@ pattern AlonzoGenesis ::
   Natural ->
   Natural ->
   Natural ->
-  AlonzoExtraConfig ->
+  Maybe AlonzoExtraConfig ->
   AlonzoGenesis
 pattern AlonzoGenesis
   { agCoinsPerUTxOWord
@@ -218,25 +221,34 @@ instance ToCBOR AlonzoGenesis where
 instance FromJSON AlonzoGenesis where
   parseJSON = Aeson.withObject "Alonzo Genesis" $ \o -> do
     agCoinsPerUTxOWord <- o .: "lovelacePerUTxOWord"
-    agPlutusV1CostModel <- parseCostModelAsArray False PlutusV1 =<< o .: "plutusV1CostModel"
+    cms <- parseCostModels False =<< o .: "costModels"
     agPrices <- o .: "executionPrices"
     agMaxTxExUnits <- o .: "maxTxExUnits"
     agMaxBlockExUnits <- o .: "maxBlockExUnits"
     agMaxValSize <- o .: "maxValueSize"
     agCollateralPercentage <- o .: "collateralPercentage"
     agMaxCollateralInputs <- o .: "maxCollateralInputs"
-    agExtraConfig <- o .: "extraConfig"
+    agExtraConfig <- o .:? "extraConfig"
+    agPlutusV1CostModel <-
+      case Map.toList (costModelsValid cms) of
+        [] -> fail "Expected \"PlutusV1\" cost model to be supplied"
+        [(PlutusV1, pv1CostModel)] -> pure pv1CostModel
+        _ ->
+          fail $
+            "Only PlutusV1 CostModel is allowed in the AlonzoGenesis, but "
+              <> List.intercalate ", " (map show . Map.keys $ flattenCostModels cms)
+              <> " were supplied. Use \"extraConfig\" if you need to inject other cost models for testing."
     return AlonzoGenesis {..}
 
 instance ToKeyValuePairs AlonzoGenesis where
   toKeyValuePairs ag =
     [ "lovelacePerUTxOWord" .= agCoinsPerUTxOWord ag
-    , "plutusV1CostModel" .= agPlutusV1CostModel ag
+    , "costModels" .= mkCostModels (Map.singleton PlutusV1 $ agPlutusV1CostModel ag)
     , "executionPrices" .= agPrices ag
     , "maxTxExUnits" .= agMaxTxExUnits ag
     , "maxBlockExUnits" .= agMaxBlockExUnits ag
     , "maxValueSize" .= agMaxValSize ag
     , "collateralPercentage" .= agCollateralPercentage ag
     , "maxCollateralInputs" .= agMaxCollateralInputs ag
-    , "extraConfig" .= agExtraConfig ag
     ]
+      ++ ["extraConfig" .= extraConfig | Just extraConfig <- [agExtraConfig ag]]

eras/alonzo/impl/src/Cardano/Ledger/Alonzo/Transition.hs

@@ -1,18 +1,23 @@
 {-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# OPTIONS_GHC -Wno-orphans #-}
 
 module Cardano.Ledger.Alonzo.Transition (
   TransitionConfig (..),
+  alonzoInjectCostModels,
 ) where
 
+import Cardano.Ledger.Alonzo.Core (AlonzoEraPParams, ppCostModelsL)
 import Cardano.Ledger.Alonzo.Era
-import Cardano.Ledger.Alonzo.Genesis (AlonzoGenesis)
+import Cardano.Ledger.Alonzo.Genesis
 import Cardano.Ledger.Alonzo.Translation ()
 import Cardano.Ledger.Mary
 import Cardano.Ledger.Mary.Transition (TransitionConfig (MaryTransitionConfig))
+import Cardano.Ledger.Plutus.CostModels (CostModels, updateCostModels)
+import Cardano.Ledger.Shelley.LedgerState
 import Cardano.Ledger.Shelley.Transition
 import GHC.Generics
 import Lens.Micro
@@ -27,7 +32,8 @@ instance EraTransition AlonzoEra where
 
   mkTransitionConfig = AlonzoTransitionConfig
 
-  injectIntoTestState = shelleyRegisterInitialFundsThenStaking
+  injectIntoTestState cfg =
+    shelleyRegisterInitialFundsThenStaking cfg . alonzoInjectCostModels cfg
 
   tcPreviousEraConfigL =
     lens atcMaryTransitionConfig (\atc pc -> atc {atcMaryTransitionConfig = pc})
@@ -36,3 +42,20 @@ instance EraTransition AlonzoEra where
     lens atcAlonzoGenesis (\atc ag -> atc {atcAlonzoGenesis = ag})
 
 instance NoThunks (TransitionConfig AlonzoEra)
+
+alonzoInjectCostModels ::
+  (EraTransition era, AlonzoEraPParams era) =>
+  TransitionConfig AlonzoEra -> NewEpochState era -> NewEpochState era
+alonzoInjectCostModels cfg =
+  case agExtraConfig $ cfg ^. tcTranslationContextL of
+    Nothing -> id
+    Just aec -> overrideCostModels (aecCostModels aec)
+
+overrideCostModels ::
+  (EraTransition era, AlonzoEraPParams era) =>
+  Maybe CostModels ->
+  NewEpochState era ->
+  NewEpochState era
+overrideCostModels = \case
+  Nothing -> id
+  Just cms -> nesEsL . curPParamsEpochStateL . ppCostModelsL %~ updateCostModels cms

eras/alonzo/impl/testlib/Test/Cardano/Ledger/Alonzo/Examples.hs

@@ -15,7 +15,7 @@ module Test.Cardano.Ledger.Alonzo.Examples (
 
 import Cardano.Ledger.Alonzo (AlonzoEra)
 import Cardano.Ledger.Alonzo.Core
-import Cardano.Ledger.Alonzo.Genesis (AlonzoExtraConfig (..), AlonzoGenesis (..))
+import Cardano.Ledger.Alonzo.Genesis (AlonzoGenesis (..))
 import Cardano.Ledger.Alonzo.Plutus.Context (EraPlutusTxInfo)
 import Cardano.Ledger.Alonzo.Scripts (
   AlonzoPlutusPurpose (..),
@@ -217,5 +217,5 @@ exampleAlonzoGenesis =
     , agMaxValSize = 1234
     , agCollateralPercentage = 20
     , agMaxCollateralInputs = 30
-    , agExtraConfig = AlonzoExtraConfig Nothing
+    , agExtraConfig = Nothing
     }

eras/alonzo/impl/testlib/Test/Cardano/Ledger/Alonzo/ImpTest.hs

@@ -47,7 +47,7 @@ module Test.Cardano.Ledger.Alonzo.ImpTest (
 import Cardano.Ledger.Address (Addr (..))
 import Cardano.Ledger.Alonzo (AlonzoEra)
 import Cardano.Ledger.Alonzo.Core
-import Cardano.Ledger.Alonzo.Genesis (AlonzoExtraConfig (..), AlonzoGenesis (..))
+import Cardano.Ledger.Alonzo.Genesis (AlonzoGenesis (..))
 import Cardano.Ledger.Alonzo.Plutus.Context (ContextError)
 import Cardano.Ledger.Alonzo.Plutus.Evaluate (
   collectPlutusScriptsWithContext,
@@ -432,7 +432,7 @@ instance ShelleyEraImp AlonzoEra where
         , agMaxValSize = 5000
         , agCollateralPercentage = 150
         , agMaxCollateralInputs = 3
-        , agExtraConfig = AlonzoExtraConfig Nothing
+        , agExtraConfig = Nothing
         }
 
   impSatisfyNativeScript = impAllegraSatisfyNativeScript