From cc056254bd25be4847fcb751da04ed07dc00ce54 Mon Sep 17 00:00:00 2001
From: effectfully <effectfully@gmail.com>
Date: Fri, 14 Jun 2024 21:36:55 +0200
Subject: [PATCH] [Evaluation] Move stuff under `enterComputeCek` (#6156)

This polishes the structure of the CEK machine code a little: moves definitions having CEK-specific constraints around (as per `Note [Compilation peculiarities]`), removes outdated comments and moves pragmas around.
---
 .../Evaluation/Machine/Cek/Internal.hs        | 130 +++++++-----------
 .../Machine/SteppableCek/Internal.hs          |  22 +--
 2 files changed, 64 insertions(+), 88 deletions(-)

diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Evaluation/Machine/Cek/Internal.hs b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Evaluation/Machine/Cek/Internal.hs
index 1826d96dbe6..f6fe8bc0456 100644
--- a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Evaluation/Machine/Cek/Internal.hs
+++ b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Evaluation/Machine/Cek/Internal.hs
@@ -1,9 +1,5 @@
 -- editorconfig-checker-disable-file
 -- | The CEK machine.
--- The CEK machine relies on variables having non-equal 'Unique's whenever they have non-equal
--- string names. I.e. 'Unique's are used instead of string names. This is for efficiency reasons.
--- The CEK machines handles name capture by design.
-
 
 {-# LANGUAGE BangPatterns             #-}
 {-# LANGUAGE ConstraintKinds          #-}
@@ -49,7 +45,6 @@ module UntypedPlutusCore.Evaluation.Machine.Cek.Internal
     , ThrowableBuiltins
     , extractEvaluationResult
     , unsafeToEvaluationResult
-    , spendBudgetStreamCek
     , runCekDeBruijn
     , dischargeCekValue
     , Context (..)
@@ -141,30 +136,24 @@ Hence we don't export 'computeCek' and instead define 'runCek' in this file and
 though the rest of the user-facing API (which 'runCek' is a part of) is defined downstream.
 
 Another problem is handling mutual recursion in the 'computeCek'/'returnCek'/'forceEvaluate'/etc
-family. If we keep these functions at the top level, GHC won't be able to pull the constraints out of
-the family (confirmed by inspecting Core: GHC thinks that since the superclass constraints
+family. If we keep these functions at the top level, GHC won't be able to pull the constraints out
+of the family (confirmed by inspecting Core: GHC thinks that since the superclass constraints
 populating the dictionary representing the @Ix fun@ constraint are redundant, they can be replaced
 with calls to 'error' in a recursive call, but that changes the dictionary and so it can no longer
 be pulled out of recursion). But that entails passing a redundant argument around, which slows down
 the machine a tiny little bit.
 
-Hence we define a number of the functions as local functions making use of a
-shared context from their parent function. This also allows GHC to inline almost
-all of the machine into a single definition (with a bunch of recursive join
-points in it).
+Hence we define a all happy-path functions having CEK-machine-specific constraints as local
+functions making use of a shared context from their parent function. This also allows GHC to inline
+almost all of the machine into a single definition (with a bunch of recursive join points in it).
 
 In general, it's advised to run benchmarks (and look at Core output if the results are suspicious)
 on any changes in this file.
 
-Finally, it's important to put bang patterns on any Int arguments to ensure that GHC unboxes them:
+Finally, it's important to put bang patterns on any 'Int' arguments to ensure that GHC unboxes them:
 this can make a surprisingly large difference.
 -}
 
-{- Note [Scoping]
-The CEK machine does not rely on the global uniqueness condition, so the renamer pass is not a
-prerequisite. The CEK machine correctly handles name shadowing.
--}
-
 -- | The 'Term's that CEK can execute must have DeBruijn binders
 -- 'Name' is not necessary but we leave it here for simplicity and debuggability.
 type NTerm uni fun = Term NamedDeBruijn uni fun
@@ -466,17 +455,6 @@ instance ThrowableBuiltins uni fun => MonadError (CekEvaluationException NamedDe
         unsafeRunCekM :: CekM uni fun s a -> IO a
         unsafeRunCekM = unsafeSTToIO . unCekM
 
--- It would be really nice to define this instance, so that we can use 'makeKnown' directly in
--- the 'CekM' monad without the 'WithEmitterT' nonsense. Unfortunately, GHC doesn't like
--- implicit params in instance contexts. As GHC's docs explain:
---
--- > Reason: exactly which implicit parameter you pick up depends on exactly where you invoke a
--- > function. But the "invocation" of instance declarations is done behind the scenes by the
--- > compiler, so it's hard to figure out exactly where it is done. The easiest thing is to outlaw
--- > the offending types.
--- instance GivenCekEmitter s => MonadEmitter (CekM uni fun s) where
---     emit = emitCek
-
 instance AsEvaluationFailure CekUserError where
     _EvaluationFailure = _EvaluationFailureVia CekEvaluationFailure
 
@@ -493,10 +471,6 @@ instance Pretty CekUserError where
           ]
     pretty CekEvaluationFailure = "The machine terminated because of an error, either from a built-in function or from an explicit use of 'error'."
 
-spendBudgetCek :: GivenCekSpender uni fun s => ExBudgetCategory fun -> ExBudget -> CekM uni fun s ()
-spendBudgetCek = let (CekBudgetSpender spend) = ?cekBudgetSpender in spend
-
--- see Note [Scoping].
 -- | Instantiate all the free variables of a term by looking them up in an environment.
 -- Mutually recursive with dischargeCekVal.
 dischargeCekValEnv :: forall uni fun ann. CekValEnv uni fun ann -> NTerm uni fun () -> NTerm uni fun ()
@@ -623,49 +597,6 @@ runCekM (MachineParameters costs runtime) (ExBudgetMode getExBudgetInfo) (Emitte
     pure (errOrRes, st, logs)
 {-# INLINE runCekM #-}
 
--- | Look up a variable name in the environment.
-lookupVarName
-    :: forall uni fun ann s
-    . ThrowableBuiltins uni fun
-    => NamedDeBruijn -> CekValEnv uni fun ann -> CekM uni fun s (CekValue uni fun ann)
-lookupVarName varName@(NamedDeBruijn _ varIx) varEnv =
-    case varEnv `Env.indexOne` coerce varIx of
-        Nothing  -> throwingWithCause _MachineError OpenTermEvaluatedMachineError $ Just var where
-            var = Var () varName
-        Just val -> pure val
-
--- | Spend each budget from the given stream of budgets.
-spendBudgetStreamCek
-    :: GivenCekReqs uni fun ann s
-    => ExBudgetCategory fun
-    -> ExBudgetStream
-    -> CekM uni fun s ()
-spendBudgetStreamCek exCat = go where
-    go (ExBudgetLast budget)         = spendBudgetCek exCat budget
-    go (ExBudgetCons budget budgets) = spendBudgetCek exCat budget *> go budgets
-{-# INLINE spendBudgetStreamCek #-}
-
--- | Take pieces of a possibly partial builtin application and either create a 'CekValue' using
--- 'makeKnown' or a partial builtin application depending on whether the built-in function is
--- fully saturated or not.
-evalBuiltinApp
-    :: (GivenCekReqs uni fun ann s, ThrowableBuiltins uni fun)
-    => fun
-    -> NTerm uni fun ()
-    -> BuiltinRuntime (CekValue uni fun ann)
-    -> CekM uni fun s (CekValue uni fun ann)
-evalBuiltinApp fun term runtime = case runtime of
-    BuiltinCostedResult budgets getX -> do
-        spendBudgetStreamCek (BBuiltinApp fun) budgets
-        case getX of
-            BuiltinSuccess x              -> pure x
-            BuiltinSuccessWithLogs logs x -> ?cekEmitter logs $> x
-            BuiltinFailure logs err       -> do
-                ?cekEmitter logs
-                throwBuiltinErrorWithCause term err
-    _ -> pure $ VBuiltin fun term runtime
-{-# INLINE evalBuiltinApp #-}
-
 -- See Note [Compilation peculiarities].
 -- | The entering point to the CEK machine's engine.
 enterComputeCek
@@ -849,16 +780,18 @@ enterComputeCek = computeCek
         let ctr = ?cekStepCounter
         iforCounter_ ctr spend
         resetCounter ctr
+    -- It's very important for this definition not to get inlined. Inlining it caused performance to
+    -- degrade by 16+%: https://github.com/IntersectMBO/plutus/pull/5931
+    {-# NOINLINE spendAccumulatedBudget #-}
 
     -- Making this a definition of its own causes it to inline better than actually writing it inline, for
     -- some reason.
     -- Skip index 7, that's the total counter!
     -- See Note [Structure of the step counter]
-    {-# INLINE spend #-}
     spend !i !w = unless (i == (fromIntegral $ natVal $ Proxy @TotalCountIndex)) $
-      let kind = toEnum i in spendBudgetCek (BStep kind) (stimes w (cekStepCost ?cekCosts kind))
+      let kind = toEnum i in spendBudget (BStep kind) (stimes w (cekStepCost ?cekCosts kind))
+    {-# INLINE spend #-}
 
-    {-# INLINE stepAndMaybeSpend #-}
     -- | Accumulate a step, and maybe spend the budget that has accumulated for a number of machine steps, but only if we've exceeded our slippage.
     stepAndMaybeSpend :: StepKind -> CekM uni fun s ()
     stepAndMaybeSpend !kind = do
@@ -873,6 +806,45 @@ enterComputeCek = computeCek
         -- There's no risk of overflow here, since we only ever increment the total
         -- steps by 1 and then check this condition.
         when (unbudgetedStepsTotal >= ?cekSlippage) spendAccumulatedBudget
+    {-# INLINE stepAndMaybeSpend #-}
+
+    -- | Take pieces of a possibly partial builtin application and either create a 'CekValue' using
+    -- 'makeKnown' or a partial builtin application depending on whether the built-in function is
+    -- fully saturated or not.
+    evalBuiltinApp
+        :: fun
+        -> NTerm uni fun ()
+        -> BuiltinRuntime (CekValue uni fun ann)
+        -> CekM uni fun s (CekValue uni fun ann)
+    evalBuiltinApp fun term runtime = case runtime of
+        BuiltinCostedResult budgets0 getX -> do
+            let exCat = BBuiltinApp fun
+                spendBudgets (ExBudgetLast budget) = spendBudget exCat budget
+                spendBudgets (ExBudgetCons budget budgets) =
+                    spendBudget exCat budget *> spendBudgets budgets
+            spendBudgets budgets0
+            case getX of
+                BuiltinSuccess x              -> pure x
+                BuiltinSuccessWithLogs logs x -> ?cekEmitter logs $> x
+                BuiltinFailure logs err       -> do
+                    ?cekEmitter logs
+                    throwBuiltinErrorWithCause term err
+        _ -> pure $ VBuiltin fun term runtime
+    {-# INLINE evalBuiltinApp #-}
+
+    spendBudget :: ExBudgetCategory fun -> ExBudget -> CekM uni fun s ()
+    spendBudget = unCekBudgetSpender ?cekBudgetSpender
+    {-# INLINE spendBudget #-}
+
+    -- | Look up a variable name in the environment.
+    lookupVarName :: NamedDeBruijn -> CekValEnv uni fun ann -> CekM uni fun s (CekValue uni fun ann)
+    lookupVarName varName@(NamedDeBruijn _ varIx) varEnv =
+        case varEnv `Env.indexOne` coerce varIx of
+            Nothing  ->
+                throwingWithCause _MachineError OpenTermEvaluatedMachineError $
+                    Just $ Var () varName
+            Just val -> pure val
+    {-# INLINE lookupVarName #-}
 
 -- See Note [Compilation peculiarities].
 -- | Evaluate a term using the CEK machine and keep track of costing, logging is optional.
@@ -885,7 +857,7 @@ runCekDeBruijn
     -> (Either (CekEvaluationException NamedDeBruijn uni fun) (NTerm uni fun ()), cost, [Text])
 runCekDeBruijn params mode emitMode term =
     runCekM params mode emitMode $ do
-        spendBudgetCek BStartup $ runIdentity $ cekStartupCost ?cekCosts
+        unCekBudgetSpender ?cekBudgetSpender BStartup $ runIdentity $ cekStartupCost ?cekCosts
         enterComputeCek NoFrame Env.empty term
 
 {- Note [Accumulators for terms]
diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Evaluation/Machine/SteppableCek/Internal.hs b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Evaluation/Machine/SteppableCek/Internal.hs
index 58e259c5977..953d1f85326 100644
--- a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Evaluation/Machine/SteppableCek/Internal.hs
+++ b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Evaluation/Machine/SteppableCek/Internal.hs
@@ -39,15 +39,14 @@ module UntypedPlutusCore.Evaluation.Machine.SteppableCek.Internal
     )
 where
 
-import Control.Monad.Primitive
 import PlutusCore.Builtin
 import PlutusCore.DeBruijn
 import PlutusCore.Evaluation.Machine.ExBudget
+import PlutusCore.Evaluation.Machine.ExBudgetStream
 import PlutusCore.Evaluation.Machine.Exception
 import PlutusCore.Evaluation.Machine.MachineParameters
 import PlutusCore.Evaluation.Result
 import PlutusPrelude
-import Universe
 import UntypedPlutusCore.Core
 import UntypedPlutusCore.Evaluation.Machine.Cek.CekMachineCosts (CekMachineCosts,
                                                                  CekMachineCostsBase (..))
@@ -57,6 +56,7 @@ import UntypedPlutusCore.Evaluation.Machine.Cek.StepCounter
 
 import Control.Lens hiding (Context)
 import Control.Monad
+import Control.Monad.Primitive
 import Data.Proxy
 import Data.RandomAccessList.Class qualified as Env
 import Data.Semigroup (stimes)
@@ -64,6 +64,7 @@ import Data.Text (Text)
 import Data.Vector qualified as V
 import Data.Word (Word64)
 import GHC.TypeNats
+import Universe
 
 {- Note [Debuggable vs Original versions of CEK]
 
@@ -283,7 +284,7 @@ runCekDeBruijn
     -> (Either (CekEvaluationException NamedDeBruijn uni fun) (NTerm uni fun ()), cost, [Text])
 runCekDeBruijn params mode emitMode term =
     runCekM params mode emitMode $ do
-        spendBudgetCek BStartup $ runIdentity $ cekStartupCost ?cekCosts
+        spendBudget BStartup $ runIdentity $ cekStartupCost ?cekCosts
         enterComputeCek NoFrame Env.empty term
 
 -- See Note [Compilation peculiarities].
@@ -442,8 +443,12 @@ evalBuiltinApp
     -> BuiltinRuntime (CekValue uni fun ann)
     -> CekM uni fun s (CekValue uni fun ann)
 evalBuiltinApp fun term runtime = case runtime of
-    BuiltinCostedResult budgets getX -> do
-        spendBudgetStreamCek (BBuiltinApp fun) budgets
+    BuiltinCostedResult budgets0 getX -> do
+        let exCat = BBuiltinApp fun
+            spendBudgets (ExBudgetLast budget) = spendBudget exCat budget
+            spendBudgets (ExBudgetCons budget budgets) =
+                spendBudget exCat budget *> spendBudgets budgets
+        spendBudgets budgets0
         case getX of
             BuiltinSuccess x              -> pure x
             BuiltinSuccessWithLogs logs x -> ?cekEmitter logs $> x
@@ -453,11 +458,10 @@ evalBuiltinApp fun term runtime = case runtime of
     _ -> pure $ VBuiltin fun term runtime
 {-# INLINE evalBuiltinApp #-}
 
-spendBudgetCek :: GivenCekSpender uni fun s => ExBudgetCategory fun -> ExBudget -> CekM uni fun s ()
-spendBudgetCek = let (CekBudgetSpender spend) = ?cekBudgetSpender in spend
+spendBudget :: GivenCekSpender uni fun s => ExBudgetCategory fun -> ExBudget -> CekM uni fun s ()
+spendBudget = unCekBudgetSpender ?cekBudgetSpender
 
 -- | Spend the budget that has been accumulated for a number of machine steps.
---
 spendAccumulatedBudget :: (GivenCekReqs uni fun ann s) => CekM uni fun s ()
 spendAccumulatedBudget = do
     let ctr = ?cekStepCounter
@@ -470,7 +474,7 @@ spendAccumulatedBudget = do
     -- See Note [Structure of the step counter]
     {-# INLINE spend #-}
     spend !i !w = unless (i == (fromIntegral $ natVal $ Proxy @TotalCountIndex)) $
-      let kind = toEnum i in spendBudgetCek (BStep kind) (stimes w (cekStepCost ?cekCosts kind))
+      let kind = toEnum i in spendBudget (BStep kind) (stimes w (cekStepCost ?cekCosts kind))
 
 -- | Accumulate a step, and maybe spend the budget that has accumulated for a number of machine steps, but only if we've exceeded our slippage.
 stepAndMaybeSpend :: (GivenCekReqs uni fun ann s) => StepKind -> CekM uni fun s ()