From bbeb1a4a5669223af51931927b7158a9715e1e87 Mon Sep 17 00:00:00 2001
From: Ziyang Liu <unsafeFixIO@gmail.com>
Date: Mon, 3 Jun 2024 14:15:55 -0700
Subject: [PATCH] Fix V2's costModelParamsForTesting (#6166)

---
 .../Evaluation/Machine/ExBudgetingDefaults.hs |  5 ++++-
 plutus-ledger-api/test/Spec.hs                | 19 ++++++-------------
 .../Test/V2/EvaluationContext.hs              |  3 ++-
 3 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/plutus-core/plutus-core/src/PlutusCore/Evaluation/Machine/ExBudgetingDefaults.hs b/plutus-core/plutus-core/src/PlutusCore/Evaluation/Machine/ExBudgetingDefaults.hs
index 148f369897c..2f4d119c813 100644
--- a/plutus-core/plutus-core/src/PlutusCore/Evaluation/Machine/ExBudgetingDefaults.hs
+++ b/plutus-core/plutus-core/src/PlutusCore/Evaluation/Machine/ExBudgetingDefaults.hs
@@ -15,6 +15,7 @@ module PlutusCore.Evaluation.Machine.ExBudgetingDefaults
     , defaultCostModelParamsForTesting
     , defaultBuiltinCostModelForTesting
     , defaultCekCostModelForTesting
+    , defaultCekCostModelForTestingB
     , unitCekMachineCosts
     , unitCekParameters
     )
@@ -219,6 +220,9 @@ defaultCostModelParamsForTesting = defaultCostModelParamsC
 defaultCekCostModelForTesting :: CostModel CekMachineCosts BuiltinCostModel
 defaultCekCostModelForTesting = cekCostModelVariantC
 
+defaultCekCostModelForTestingB :: CostModel CekMachineCosts BuiltinCostModel
+defaultCekCostModelForTestingB = cekCostModelVariantB
+
 {- A cost model with unit costs, so we can count how often each builtin is called.
   This currently works for all semantics variants because to date we have only
   ever added new builtins and never removed any. -}
@@ -333,4 +337,3 @@ unitCekParameters =
     -- See Note [noinline for saving on ticks].
     noinline mkMachineParameters def $
         CostModel unitCekMachineCosts unitCostBuiltinCostModel
-
diff --git a/plutus-ledger-api/test/Spec.hs b/plutus-ledger-api/test/Spec.hs
index f1eb5d58551..8e312a47b30 100644
--- a/plutus-ledger-api/test/Spec.hs
+++ b/plutus-ledger-api/test/Spec.hs
@@ -19,7 +19,7 @@ import Spec.Versions qualified
 
 import Test.Tasty
 import Test.Tasty.HUnit
--- import Test.Tasty.QuickCheck
+import Test.Tasty.QuickCheck
 
 import Control.Monad.Writer
 import Data.Int (Int64)
@@ -36,10 +36,6 @@ v1_evalCtxForTesting = fst $ unsafeFromRight $ runWriterT $ V1.mkEvaluationConte
 v3_evalCtxTooFewParams :: V3.EvaluationContext
 v3_evalCtxTooFewParams = fst $ unsafeFromRight $ runWriterT $ V3.mkEvaluationContext (take 223 $ fmap snd V3.costModelParamsForTesting)
 
-
--- ** FIXME: the change in the structure of the cost models has invalidated a number of the
--- plutus-ledger-api tests.  We need to work out how to fix these properly.
-{-
 alwaysTrue :: TestTree
 alwaysTrue = testCase "always true script returns true" $
     let script = either (error . show) id $ V1.deserialiseScript alonzoPV (alwaysSucceedingNAryFunction 2)
@@ -51,7 +47,6 @@ alwaysFalse = testCase "always false script returns false" $
     let script = either (error . show) id $ V1.deserialiseScript alonzoPV (alwaysFailingNAryFunction 2)
         (_, res) = V1.evaluateScriptCounting alonzoPV V1.Quiet v1_evalCtxForTesting script [I 1, I 2]
     in assertBool "fails" (isLeft res)
--}
 
 unavailableBuiltins :: TestTree
 unavailableBuiltins = testCase "builtins are unavailable before Alonzo" $
@@ -71,7 +66,6 @@ integerToByteStringExceedsBudget = testCase "integerToByteString should exceed b
         Left _ -> assertFailure "fails"
         Right (ExBudget cpu _mem) -> assertBool "did not exceed budget" (cpu >= fromIntegral (maxBound :: Int64))
 
-{-  ** FIXME: These don't work with the new cost model setup
 saltedFunction :: TestTree
 saltedFunction =
     let evaluate ss ss' args =
@@ -112,16 +106,15 @@ saltedFunction =
             f'' = saltFunction salt' f
         in salt /= salt' ==> f' /= f''
     ]
--}
+
 
 tests :: TestTree
 tests = testGroup "plutus-ledger-api"[
     testGroup "basic evaluation tests" [
-{-          alwaysTrue
-          , alwaysFalse
-          , saltedFunction
--}
-          unavailableBuiltins
+          alwaysTrue
+        , alwaysFalse
+        , saltedFunction
+        , unavailableBuiltins
         , availableBuiltins
         , integerToByteStringExceedsBudget
     ]
diff --git a/plutus-ledger-api/testlib/PlutusLedgerApi/Test/V2/EvaluationContext.hs b/plutus-ledger-api/testlib/PlutusLedgerApi/Test/V2/EvaluationContext.hs
index e7bbec6cd2d..932486ee8d3 100644
--- a/plutus-ledger-api/testlib/PlutusLedgerApi/Test/V2/EvaluationContext.hs
+++ b/plutus-ledger-api/testlib/PlutusLedgerApi/Test/V2/EvaluationContext.hs
@@ -9,6 +9,7 @@ module PlutusLedgerApi.Test.V2.EvaluationContext
     ) where
 
 import PlutusCore.Evaluation.Machine.BuiltinCostModel
+import PlutusCore.Evaluation.Machine.ExBudgetingDefaults
 import PlutusCore.Evaluation.Machine.MachineParameters
 import PlutusLedgerApi.Test.Common.EvaluationContext as Common
 import PlutusLedgerApi.Test.V3.EvaluationContext qualified as V3
@@ -27,7 +28,7 @@ costModelParamsForTesting = Map.toList $ fromJust $
 
 -- | The PlutusV2 "cost model" is constructed by the v3 "cost model", by clearing v3 introductions.
 mCostModel :: MCostModel
-mCostModel = V3.mCostModel
+mCostModel = toMCostModel defaultCekCostModelForTestingB
            & machineCostModel
            %~ V3.clearMachineCostModel
            & builtinCostModel