Sam/expose budget multiplier

README.md

@@ -126,6 +126,8 @@ Top-level configuration file fields:
     collectStats: `true` or `false`
         Save some stats during contract run (only transactions execution
         budgets supported atm) (default: false)
+    budgetMultiplier: rational multiplier in form `1` or `1 % 2`
+         (default: 1)
 ```
 
 To run the fake PAB, you need to prepare a few more things:

src/BotPlutusInterface/Config.hs

@@ -33,6 +33,7 @@ import Data.String.ToString (toString)
 import Data.Text (Text)
 import Data.Text qualified as Text
 import PlutusConfig.Base (
+  customRationalSpec,
   enumToAtom,
   filepathSpec,
   maybeSpec,
@@ -95,6 +96,7 @@ instance ToValue PABConfig where
         pcPort
         pcEnableTxEndpoint
         pcCollectStats
+        pcBudgetMultiplier
       ) =
       Sections
         ()
@@ -116,6 +118,7 @@ instance ToValue PABConfig where
         , Section () "port"               $ toValue pcPort
         , Section () "enableTxEndpoint"   $ toValue pcEnableTxEndpoint
         , Section () "collectStats"       $ toValue pcCollectStats
+        , Section () "budgetMultiplier"   $ toValue pcBudgetMultiplier
         ]
 {- ORMOLU_ENABLE -}
 
@@ -206,6 +209,13 @@ pabConfigSpec = sectionsSpec "PABConfig" $ do
       trueOrFalseSpec
       "Save some stats during contract run (only transactions execution budgets supported atm)"
 
+  pcBudgetMultiplier <-
+    sectionWithDefault'
+      (pcBudgetMultiplier def)
+      "budgetMultiplier"
+      customRationalSpec
+      "Multiplier on the budgets automatically calculated"
+
   pure PABConfig {..}
 
 docPABConfig :: String

src/BotPlutusInterface/ExBudget.hs

@@ -14,11 +14,15 @@ import BotPlutusInterface.Types (
   TxFile (..),
  )
 import Cardano.Api qualified as CAPI
+import Cardano.Api.Shelley (ProtocolParameters (protocolParamMaxTxExUnits))
+import Cardano.Prelude (maybeToEither)
 import Control.Arrow (left)
+import Data.Either (rights)
 import Data.List (sort)
 import Data.Map (Map)
 import Data.Map qualified as Map
 import Data.Text qualified as Text
+import GHC.Natural (Natural)
 import Ledger (ExBudget (ExBudget), ExCPU (ExCPU), ExMemory (ExMemory), MintingPolicyHash, TxOutRef)
 import Ledger.Tx.CardanoAPI (fromCardanoPolicyId, fromCardanoTxIn)
 import System.Directory.Internal.Prelude (getEnv)
@@ -43,13 +47,47 @@ estimateBudget pabConf txFile = do
       (getExUnits debugNodeInf)
       txBody
 
-  let txBudget = do
-        body <- txBody
-        budget <- budgetRes
-        (spendingBudgets, policyBudgets) <- mkBudgetMaps budget body
-        Right $ TxBudget spendingBudgets policyBudgets
+  return $
+    do
+      body <- txBody
+      budget <- budgetRes
+      maxUnits <- maybeToEither (BudgetEstimationError "Missing max units in parameters") $ protocolParamMaxTxExUnits pabConf.pcProtocolParams
+
+      scaledBudget <- getScaledBudget maxUnits pabConf.pcBudgetMultiplier budget
+
+      (spendingBudgets, policyBudgets) <- mkBudgetMaps scaledBudget body
+
+      Right $ TxBudget spendingBudgets policyBudgets
+
+-- | Scale the budget clamping the total to the parameter limits
+getScaledBudget :: CAPI.ExecutionUnits -> Rational -> ExUnitsMap -> Either BudgetEstimationError ExUnitsMap
+getScaledBudget maxUnits scaler budget =
+  if fst scalers >= 1 && snd scalers >= 1
+    then Right $ fmap (fmap $ scaleBudget scalers) budget
+    else
+      Left $
+        BudgetEstimationError $
+          Text.pack $
+            "Exceeded global transaction budget\nCalculated: " ++ show budgetSum ++ "\nLimit: " ++ show maxUnits
+  where
+    budgetSum = foldr addBudgets (CAPI.ExecutionUnits 0 0) $ rights $ Map.elems budget
+    scalers =
+      ( clampedScaler (CAPI.executionSteps budgetSum) (CAPI.executionSteps maxUnits) scaler
+      , clampedScaler (CAPI.executionMemory budgetSum) (CAPI.executionMemory maxUnits) scaler
+      )
+
+clampedScaler :: Natural -> Natural -> Rational -> Rational
+clampedScaler 0 _ scaler = scaler
+clampedScaler val maxVal scaler = min scaler (toRational maxVal / toRational val)
+
+-- | Scale the budget by the multipliers in config
+scaleBudget :: (Rational, Rational) -> CAPI.ExecutionUnits -> CAPI.ExecutionUnits
+scaleBudget (stepsScaler, memScaler) (CAPI.ExecutionUnits steps mem) = CAPI.ExecutionUnits (scale steps stepsScaler) (scale mem memScaler)
+  where
+    scale x scaler = round $ toRational x * scaler
 
-  return txBudget
+addBudgets :: CAPI.ExecutionUnits -> CAPI.ExecutionUnits -> CAPI.ExecutionUnits
+addBudgets (CAPI.ExecutionUnits steps mem) (CAPI.ExecutionUnits steps' mem') = CAPI.ExecutionUnits (steps + steps') (mem + mem')
 
 -- | Deserialize transaction body from ".signed" file
 deserialiseSigned :: FilePath -> IO (Either BudgetEstimationError (CAPI.Tx CAPI.AlonzoEra))

src/BotPlutusInterface/QueryNode.hs

@@ -3,6 +3,7 @@
 -- | Several query functions to query local node
 module BotPlutusInterface.QueryNode (
   NodeInfo (..),
+  NodeQueryError (..),
   queryProtocolParams,
   querySystemStart,
   queryEraHistory,

src/BotPlutusInterface/Types.hs

@@ -83,6 +83,7 @@ data PABConfig = PABConfig
   , pcPort :: !Port
   , pcEnableTxEndpoint :: !Bool
   , pcCollectStats :: !Bool
+  , pcBudgetMultiplier :: !Rational
   }
   deriving stock (Show, Eq)
 
@@ -221,6 +222,7 @@ instance Default PABConfig where
       , pcPort = 9080
       , pcEnableTxEndpoint = False
       , pcCollectStats = False
+      , pcBudgetMultiplier = 1
       }
 
 data RawTx = RawTx

test/Spec/BotPlutusInterface/Config.hs

@@ -113,4 +113,5 @@ pabConfigExample =
     , pcPort = 1021
     , pcEnableTxEndpoint = True
     , pcCollectStats = False
+    , pcBudgetMultiplier = 1
     }