Add a debug flag for uplc eval

Cargo.lock

@@ -75,6 +75,7 @@ dependencies = [
  "rand",
  "regex",
  "serde_json",
+ "strum",
  "thiserror 1.0.69",
  "uplc",
  "xdg",
@@ -592,7 +593,7 @@ version = "4.5.18"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4ac6a0c7b1a9e9a5186361f67dfa1b88213572f427fb9ab038efb2bd8c582dab"
 dependencies = [
- "heck 0.5.0",
+ "heck",
  "proc-macro2",
  "quote",
  "syn 2.0.91",
@@ -1332,12 +1333,6 @@ version = "0.15.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289"
 
-[[package]]
-name = "heck"
-version = "0.4.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8"
-
 [[package]]
 name = "heck"
 version = "0.5.0"
@@ -3162,21 +3157,24 @@ checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f"
 
 [[package]]
 name = "strum"
-version = "0.24.1"
+version = "0.26.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "063e6045c0e62079840579a7e47a355ae92f60eb74daaf156fb1e84ba164e63f"
+checksum = "8fec0f0aef304996cf250b31b5a10dee7980c85da9d759361292b8bca5a18f06"
+dependencies = [
+ "strum_macros",
+]
 
 [[package]]
 name = "strum_macros"
-version = "0.24.3"
+version = "0.26.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1e385be0d24f186b4ce2f9982191e7101bb737312ad61c1f2f984f34bcf85d59"
+checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be"
 dependencies = [
- "heck 0.4.1",
+ "heck",
  "proc-macro2",
  "quote",
  "rustversion",
- "syn 1.0.109",
+ "syn 2.0.91",
 ]
 
 [[package]]
@@ -3613,7 +3611,6 @@ dependencies = [
  "serde",
  "serde_json",
  "strum",
- "strum_macros",
  "thiserror 1.0.69",
  "walkdir",
 ]

Cargo.toml

@@ -61,9 +61,6 @@ pallas-traverse = "0.31.0"
 [profile.dev.package.insta]
 opt-level = 3
 
-[profile.dev.package.similar]
-opt-level = 3
-
 # The profile that 'cargo dist' will build with
 [profile.dist]
 inherits = "release"

crates/aiken-lang/Cargo.toml

@@ -30,7 +30,7 @@ patricia_tree = "0.8.0"
 petgraph = "0.6.3"
 pretty = "0.12.3"
 serde = { version = "1.0.197", features = ["derive", "rc"] }
-strum = "0.24.1"
+strum = { version = "0.26.3", features = ["derive"] }
 thiserror = "1.0.39"
 uplc = { path = '../uplc', version = "1.1.10" }
 vec1 = "1.10.1"

crates/aiken-lang/src/builtins.rs

@@ -12,9 +12,11 @@ use crate::{
     },
     IdGenerator,
 };
+
 use indexmap::IndexMap;
 use std::{collections::HashMap, rc::Rc};
 use strum::IntoEnumIterator;
+
 use uplc::{
     builder::{CONSTR_FIELDS_EXPOSER, CONSTR_INDEX_EXPOSER},
     builtins::DefaultFunction,

crates/aiken/Cargo.toml

@@ -45,6 +45,7 @@ pallas-traverse.workspace = true
 rand = "0.8.5"
 regex = "1.7.1"
 serde_json = "1.0.94"
+strum = { version = "0.26.3", features = ["derive"] }
 thiserror = "1.0.39"
 uplc = { path = '../uplc', version = "1.1.10" }
 

crates/aiken/src/cmd/uplc/eval.rs

@@ -1,9 +1,17 @@
 use miette::IntoDiagnostic;
+use pallas_primitives::conway::Language;
 use serde_json::json;
 use std::{path::PathBuf, process};
+
+use strum::IntoEnumIterator;
+
 use uplc::{
     ast::{FakeNamedDeBruijn, Name, NamedDeBruijn, Program, Term},
-    machine::cost_model::ExBudget,
+    builtins::DefaultFunction,
+    machine::{
+        cost_model::{ExBudget, StepKind},
+        TERM_COUNT,
+    },
     parser,
 };
 
@@ -18,6 +26,9 @@ pub struct Args {
     #[clap(short, long)]
     cbor: bool,
 
+    #[clap(short, long, default_value_t = false)]
+    debug: bool,
+
     /// Arguments to pass to the UPLC program
     args: Vec<String>,
 }
@@ -27,6 +38,7 @@ pub fn exec(
         script,
         flat,
         args,
+        debug,
         cbor,
     }: Args,
 ) -> miette::Result<()> {
@@ -65,7 +77,11 @@ pub fn exec(
 
     let program = Program::<NamedDeBruijn>::try_from(program).into_diagnostic()?;
 
-    let mut eval_result = program.eval(budget);
+    let mut eval_result = if debug {
+        program.eval_debug(ExBudget::default(), &Language::PlutusV3)
+    } else {
+        program.eval(budget)
+    };
 
     let cost = eval_result.cost();
     let logs = eval_result.logs();
@@ -85,6 +101,43 @@ pub fn exec(
                 serde_json::to_string_pretty(&output).into_diagnostic()?
             );
 
+            if debug {
+                println!("---------------DEBUG------------------");
+                let costs = eval_result.debug_cost().unwrap();
+
+                let mut output = json!([]);
+                for step in StepKind::iter() {
+                    if matches!(step, StepKind::StartUp) {
+                        continue;
+                    }
+                    let i = step as usize * 2;
+
+                    if costs[i + 1] != 0 || costs[i] != 0 {
+                        output.as_array_mut().unwrap().push(json!({
+                            "step": step.to_string(),
+                            "cpu": costs[i+1],
+                            "mem": costs[i],
+                        }));
+                    }
+                }
+
+                for fun in DefaultFunction::iter() {
+                    let i = (fun as usize + TERM_COUNT) * 2;
+                    if costs[i + 1] != 0 || costs[i] != 0 {
+                        output.as_array_mut().unwrap().push(json!({
+                            "fun": fun.to_string(),
+                            "cpu": costs[i+1],
+                            "mem": costs[i],
+                        }));
+                    }
+                }
+
+                println!(
+                    "{}",
+                    serde_json::to_string_pretty(&output).into_diagnostic()?
+                );
+            }
+
             Ok(())
         }
         Err(err) => {

crates/uplc/Cargo.toml

@@ -28,8 +28,7 @@ peg = "0.8.1"
 pretty = "0.11.3"
 serde = { version = "1.0.152", features = ["derive"] }
 serde_json = "1.0.94"
-strum = "0.24.1"
+strum = { version = "0.26.3", features = ["derive"] }
-strum_macros = "0.24.3"
 thiserror = "1.0.39"
 blst = "0.3.11"
 once_cell = "1.18.0"

crates/uplc/src/ast.rs

@@ -844,7 +844,7 @@ impl From<Term<FakeNamedDeBruijn>> for Term<NamedDeBruijn> {
 impl Program<NamedDeBruijn> {
     pub fn eval(self, initial_budget: ExBudget) -> EvalResult {
         let mut machine = Machine::new(
-            Language::PlutusV2,
+            Language::PlutusV3,
             CostModel::default(),
             initial_budget,
             200,
@@ -852,7 +852,13 @@ impl Program<NamedDeBruijn> {
 
         let term = machine.run(self.term);
 
-        EvalResult::new(term, machine.ex_budget, initial_budget, machine.logs)
+        EvalResult::new(
+            term,
+            machine.ex_budget,
+            initial_budget,
+            machine.logs,
+            machine.spend_counter.map(|i| i.into()),
+        )
     }
 
     /// Evaluate a Program as a specific PlutusVersion
@@ -861,7 +867,13 @@ impl Program<NamedDeBruijn> {
 
         let term = machine.run(self.term);
 
-        EvalResult::new(term, machine.ex_budget, initial_budget, machine.logs)
+        EvalResult::new(
+            term,
+            machine.ex_budget,
+            initial_budget,
+            machine.logs,
+            machine.spend_counter.map(|i| i.into()),
+        )
     }
 
     pub fn eval_as(
@@ -881,7 +893,32 @@ impl Program<NamedDeBruijn> {
 
         let term = machine.run(self.term);
 
-        EvalResult::new(term, machine.ex_budget, budget, machine.logs)
+        EvalResult::new(
+            term,
+            machine.ex_budget,
+            budget,
+            machine.logs,
+            machine.spend_counter.map(|i| i.into()),
+        )
+    }
+
+    pub fn eval_debug(self, initial_budget: ExBudget, version: &Language) -> EvalResult {
+        let mut machine = Machine::new_debug(
+            version.clone(),
+            CostModel::default(),
+            initial_budget,
+            200, //slippage
+        );
+
+        let term = machine.run(self.term);
+
+        EvalResult::new(
+            term,
+            machine.ex_budget,
+            initial_budget,
+            machine.logs,
+            machine.spend_counter.map(|i| i.into()),
+        )
     }
 }
 

crates/uplc/src/builtins.rs

@@ -1,7 +1,7 @@
 use crate::ast::Term;
 use pallas_codec::flat::de;
 use std::{fmt::Display, rc::Rc, str::FromStr};
-use strum_macros::EnumIter;
+use strum::EnumIter;
 
 /// All the possible builtin functions in Untyped Plutus Core.
 #[repr(u8)]

crates/uplc/src/machine.rs

@@ -42,11 +42,15 @@ enum Context {
     NoFrame,
 }
 
+pub const TERM_COUNT: usize = 9;
+pub const BUILTIN_COUNT: usize = 87;
+
 pub struct Machine {
     costs: CostModel,
     pub ex_budget: ExBudget,
     slippage: u32,
     unbudgeted_steps: [u32; 10],
+    pub spend_counter: Option<[i64; (TERM_COUNT + BUILTIN_COUNT) * 2]>,
     pub logs: Vec<String>,
     version: Language,
 }
@@ -63,6 +67,24 @@ impl Machine {
             ex_budget: initial_budget,
             slippage,
             unbudgeted_steps: [0; 10],
+            spend_counter: None,
+            logs: vec![],
+            version,
+        }
+    }
+
+    pub fn new_debug(
+        version: Language,
+        costs: CostModel,
+        initial_budget: ExBudget,
+        slippage: u32,
+    ) -> Machine {
+        Machine {
+            costs,
+            ex_budget: initial_budget,
+            slippage,
+            unbudgeted_steps: [0; 10],
+            spend_counter: Some([0; (TERM_COUNT + BUILTIN_COUNT) * 2]),
             logs: vec![],
             version,
         }
@@ -324,6 +346,13 @@ impl Machine {
 
         self.spend_budget(cost)?;
 
+        if let Some(counter) = &mut self.spend_counter {
+            let i = (runtime.fun as usize + TERM_COUNT) * 2;
+
+            counter[i] += cost.mem;
+            counter[i + 1] += cost.cpu;
+        }
+
         runtime.call(&self.version, &mut self.logs)
     }
 
@@ -355,6 +384,11 @@ impl Machine {
             self.spend_budget(unspent_step_budget)?;
 
             self.unbudgeted_steps[i] = 0;
+
+            if let Some(counter) = &mut self.spend_counter {
+                counter[i * 2] += unspent_step_budget.mem;
+                counter[i * 2 + 1] += unspent_step_budget.cpu;
+            }
         }
 
         self.unbudgeted_steps[9] = 0;

crates/uplc/src/machine/cost_model.rs

@@ -4,6 +4,8 @@ use num_traits::Signed;
 use pallas_primitives::conway::Language;
 use std::collections::HashMap;
 
+use strum::{Display, EnumIter};
+
 macro_rules! hashmap {
     // map-like
     ($($k:expr => $v:expr),* $(,)?) => {{
@@ -5345,6 +5347,7 @@ pub struct TwoArgumentsQuadraticFunction {
 }
 
 #[repr(u8)]
+#[derive(Debug, EnumIter, Display, Clone, Copy)]
 pub enum StepKind {
     Constant = 0,
     Var = 1,

crates/uplc/src/machine/eval_result.rs

@@ -7,6 +7,7 @@ pub struct EvalResult {
     remaining_budget: ExBudget,
     initial_budget: ExBudget,
     logs: Vec<String>,
+    debug_cost: Option<Vec<i64>>,
 }
 
 impl EvalResult {
@@ -15,12 +16,14 @@ impl EvalResult {
         remaining_budget: ExBudget,
         initial_budget: ExBudget,
         logs: Vec<String>,
+        debug_cost: Option<Vec<i64>>,
     ) -> EvalResult {
         EvalResult {
             result,
             remaining_budget,
             initial_budget,
             logs,
+            debug_cost,
         }
     }
 
@@ -47,6 +50,10 @@ impl EvalResult {
         }
     }
 
+    pub fn debug_cost(&self) -> Option<Vec<i64>> {
+        self.debug_cost.clone()
+    }
+
     #[allow(clippy::result_unit_err)]
     pub fn unwrap_constant(self) -> Result<Constant, ()> {
         match self.result {

crates/uplc/src/machine/runtime.rs

@@ -53,7 +53,7 @@ impl From<&Language> for BuiltinSemantics {
 #[derive(Clone, Debug, PartialEq)]
 pub struct BuiltinRuntime {
     pub(super) args: Vec<Value>,
-    fun: DefaultFunction,
+    pub fun: DefaultFunction,
     pub(super) forces: u32,
 }