Allow to set initial budget

crates/uplc/src/ast.rs

@@ -519,15 +519,21 @@ impl Program<NamedDeBruijn> {
         &self,
         version: &Language,
         costs: &[i64],
+        initial_budget: Option<&ExBudget>,
     ) -> (
         Result<Term<NamedDeBruijn>, crate::machine::Error>,
         ExBudget,
         Vec<String>,
     ) {
+        let budget = match initial_budget {
+            Some(b) => b.clone(),
+            None => ExBudget::default(),
+        };
+
         let mut machine = Machine::new(
             version.clone(),
             initialize_cost_model(version, costs),
-            ExBudget::default(),
+            budget,
             200, //slippage
         );
 

crates/uplc/src/tx.rs

@@ -9,6 +9,8 @@ pub use eval::get_script_and_datum_lookup_table;
 pub use phase_one::eval_phase_one;
 use script_context::{ResolvedInput, SlotConfig};
 
+use crate::machine::cost_model::ExBudget;
+
 pub mod error;
 mod eval;
 mod phase_one;
@@ -26,6 +28,7 @@ pub fn eval_phase_two(
     tx: &MintedTx,
     utxos: &[ResolvedInput],
     cost_mdls: Option<&CostMdls>,
+    initial_budget: Option<&ExBudget>,
     slot_config: &SlotConfig,
     run_phase_one: bool,
 ) -> Result<Vec<Redeemer>, Error> {
@@ -50,6 +53,7 @@ pub fn eval_phase_two(
                     redeemer,
                     &lookup_table,
                     cost_mdls,
+                    initial_budget,
                 )?;
 
                 collected_redeemers.push(redeemer)
@@ -63,10 +67,12 @@ pub fn eval_phase_two(
 
 /// This function is the same as [`eval_phase_two`]
 /// but the inputs are raw bytes.
+/// initial_budget expects (cpu, mem).
 pub fn eval_phase_two_raw(
     tx_bytes: &[u8],
     utxos_bytes: &[(Vec<u8>, Vec<u8>)],
     cost_mdls_bytes: &[u8],
+    initial_budget: (u64, u64),
     slot_config: (u64, u64),
     run_phase_one: bool,
 ) -> Result<Vec<Vec<u8>>, Error> {
@@ -75,6 +81,11 @@ pub fn eval_phase_two_raw(
 
     let cost_mdls = CostMdls::decode_fragment(cost_mdls_bytes)?;
 
+    let budget = ExBudget {
+        cpu: initial_budget.0 as i64,
+        mem: initial_budget.1 as i64,
+    };
+
     let mut utxos = Vec::new();
 
     for (input, output) in utxos_bytes {
@@ -91,7 +102,14 @@ pub fn eval_phase_two_raw(
 
     match multi_era_tx {
         MultiEraTx::Babbage(tx) => {
-            match eval_phase_two(&tx, &utxos, Some(&cost_mdls), &sc, run_phase_one) {
+            match eval_phase_two(
+                &tx,
+                &utxos,
+                Some(&cost_mdls),
+                Some(&budget),
+                &sc,
+                run_phase_one,
+            ) {
                 Ok(redeemers) => Ok(redeemers
                     .iter()
                     .map(|r| r.encode_fragment().unwrap())

crates/uplc/src/tx/eval.rs

@@ -575,6 +575,7 @@ pub fn eval_redeemer(
     redeemer: &Redeemer,
     lookup_table: &DataLookupTable,
     cost_mdls_opt: Option<&CostMdls>,
+    initial_budget: Option<&ExBudget>,
 ) -> Result<Redeemer, Error> {
     let purpose = get_script_purpose(
         redeemer,
@@ -612,7 +613,7 @@ pub fn eval_redeemer(
                         return Err(Error::V1CostModelNotFound);
                     };
 
-                    program.eval_as(&Language::PlutusV1, costs)
+                    program.eval_as(&Language::PlutusV1, costs, initial_budget)
                 } else {
                     program.eval_v1()
                 };
@@ -658,7 +659,7 @@ pub fn eval_redeemer(
                         return Err(Error::V2CostModelNotFound);
                     };
 
-                    program.eval_as(&Language::PlutusV2, costs)
+                    program.eval_as(&Language::PlutusV2, costs, initial_budget)
                 } else {
                     program.eval()
                 };
@@ -706,7 +707,7 @@ pub fn eval_redeemer(
                         return Err(Error::V1CostModelNotFound);
                     };
 
-                    program.eval_as(&Language::PlutusV1, costs)
+                    program.eval_as(&Language::PlutusV1, costs, initial_budget)
                 } else {
                     program.eval_v1()
                 };
@@ -751,7 +752,7 @@ pub fn eval_redeemer(
                         return Err(Error::V2CostModelNotFound);
                     };
 
-                    program.eval_as(&Language::PlutusV2, costs)
+                    program.eval_as(&Language::PlutusV2, costs, initial_budget)
                 } else {
                     program.eval()
                 };

crates/uplc/src/tx/tests.rs

@@ -5,6 +5,8 @@ use pallas_primitives::{
 };
 use pallas_traverse::{Era, MultiEraTx};
 
+use crate::machine::cost_model::ExBudget;
+
 use super::{eval_phase_two, ResolvedInput, SlotConfig};
 
 #[test]
@@ -226,13 +228,25 @@ fn test_eval() {
         plutus_v2: Some(costs),
     };
 
+    let initial_budget = ExBudget {
+        cpu: 10000000000,
+        mem: 14000000,
+    };
+
     let multi_era_tx = MultiEraTx::decode(Era::Babbage, &tx_bytes)
         .or_else(|_| MultiEraTx::decode(Era::Alonzo, &tx_bytes))
         .unwrap();
     match multi_era_tx {
         MultiEraTx::Babbage(tx) => {
-            let redeemers =
+            let redeemers = eval_phase_two(
-                eval_phase_two(&tx, &utxos, Some(&cost_mdl), &slot_config, false).unwrap();
+                &tx,
+                &utxos,
+                Some(&cost_mdl),
+                Some(&initial_budget),
+                &slot_config,
+                false,
+            )
+            .unwrap();
 
             assert_eq!(redeemers.len(), 1)
         }
@@ -461,13 +475,25 @@ fn test_eval_1() {
         plutus_v2: Some(costs),
     };
 
+    let initial_budget = ExBudget {
+        cpu: 10000000000,
+        mem: 14000000,
+    };
+
     let multi_era_tx = MultiEraTx::decode(Era::Babbage, &tx_bytes)
         .or_else(|_| MultiEraTx::decode(Era::Alonzo, &tx_bytes))
         .unwrap();
     match multi_era_tx {
         MultiEraTx::Babbage(tx) => {
-            let redeemers =
+            let redeemers = eval_phase_two(
-                eval_phase_two(&tx, &utxos, Some(&cost_mdl), &slot_config, false).unwrap();
+                &tx,
+                &utxos,
+                Some(&cost_mdl),
+                Some(&initial_budget),
+                &slot_config,
+                false,
+            )
+            .unwrap();
 
             println!("{:?}", redeemers.len());
         }
@@ -531,13 +557,25 @@ fn test_eval_2() {
         plutus_v2: None,
     };
 
+    let initial_budget = ExBudget {
+        cpu: 10000000000,
+        mem: 14000000,
+    };
+
     let multi_era_tx = MultiEraTx::decode(Era::Babbage, &tx_bytes)
         .or_else(|_| MultiEraTx::decode(Era::Alonzo, &tx_bytes))
         .unwrap();
     match multi_era_tx {
         MultiEraTx::Babbage(tx) => {
-            let redeemers =
+            let redeemers = eval_phase_two(
-                eval_phase_two(&tx, &utxos, Some(&cost_mdl), &slot_config, false).unwrap();
+                &tx,
+                &utxos,
+                Some(&cost_mdl),
+                Some(&initial_budget),
+                &slot_config,
+                false,
+            )
+            .unwrap();
 
             println!("{:?}", redeemers.len());
         }
@@ -598,13 +636,26 @@ fn eval_missing_redeemer() {
         plutus_v2: None,
     };
 
+    let initial_budget = ExBudget {
+        cpu: 10000000000,
+        mem: 14000000,
+    };
+
     let multi_era_tx = MultiEraTx::decode(Era::Babbage, &tx_bytes)
         .or_else(|_| MultiEraTx::decode(Era::Alonzo, &tx_bytes))
         .unwrap();
 
     match multi_era_tx {
         MultiEraTx::Babbage(tx) => {
-            eval_phase_two(&tx, &utxos, Some(&cost_mdl), &slot_config, false).unwrap();
+            eval_phase_two(
+                &tx,
+                &utxos,
+                Some(&cost_mdl),
+                Some(&initial_budget),
+                &slot_config,
+                false,
+            )
+            .unwrap();
         }
         _ => unreachable!(),
     };
@@ -663,13 +714,26 @@ fn eval_extraneous_redeemer() {
         plutus_v2: None,
     };
 
+    let initial_budget = ExBudget {
+        cpu: 10000000000,
+        mem: 14000000,
+    };
+
     let multi_era_tx = MultiEraTx::decode(Era::Babbage, &tx_bytes)
         .or_else(|_| MultiEraTx::decode(Era::Alonzo, &tx_bytes))
         .unwrap();
 
     match multi_era_tx {
         MultiEraTx::Babbage(tx) => {
-            eval_phase_two(&tx, &utxos, Some(&cost_mdl), &slot_config, false).unwrap();
+            eval_phase_two(
+                &tx,
+                &utxos,
+                Some(&cost_mdl),
+                Some(&initial_budget),
+                &slot_config,
+                false,
+            )
+            .unwrap();
         }
         _ => unreachable!(),
     };