[gas profiler] various improments

aptos-move/aptos-gas-profiling/src/aggregate.rs

@@ -6,21 +6,14 @@ use crate::{
     log::{ExecutionAndIOCosts, ExecutionGasEvent},
     render::{Render, TableKey},
 };
-use aptos_gas_algebra::{GasQuantity, GasScalingFactor, InternalGas, InternalGasUnit};
+use aptos_gas_algebra::{GasQuantity, InternalGas, InternalGasUnit};
 use std::collections::{btree_map, BTreeMap};
 
 /// Represents an aggregation of execution gas events, including the count and total gas costs for each type of event.
 ///
 /// The events are sorted by the amount of gas used, from high to low.
 #[derive(Debug)]
 pub struct AggregatedExecutionGasEvents {
-    /// The gas scaling factor.
-    /// This is included so to make this struct self-contained, suitable for displaying in (external) gas units.
-    pub gas_scaling_factor: GasScalingFactor,
-
-    /// The total gas cost.
-    pub total: InternalGas,
-
     // TODO: Make this more strongly typed?
     pub ops: Vec<(String, usize, InternalGas)>,
     pub methods: Vec<(String, usize, InternalGas)>,
@@ -136,20 +129,22 @@ impl ExecutionAndIOCosts {
         compute_method_self_cost(&tree, &mut methods_self);
 
         AggregatedExecutionGasEvents {
-            gas_scaling_factor: self.gas_scaling_factor,
-            total: self.total,
-
             ops: into_sorted_vec(ops),
             methods: into_sorted_vec(methods),
             methods_self: into_sorted_vec(methods_self),
-            transaction_write: self.transaction_transient.unwrap_or_else(|| 0.into()),
+            transaction_write: self.transaction_transient.unwrap_or_else(InternalGas::zero),
             event_writes: into_sorted_vec(event_writes),
             storage_reads: into_sorted_vec(storage_reads),
             storage_writes: into_sorted_vec(storage_writes),
         }
     }
 }
 
+/// Check if a name looks like a function call (e.g., "0xcafe::module::function")
+fn is_function_name(name: &str) -> bool {
+    name.contains("::")
+}
+
 fn compute_method_recursive_cost(
     node: &Node<GasQuantity<InternalGasUnit>>,
     methods: &mut BTreeMap<String, (usize, GasQuantity<InternalGasUnit>)>,
@@ -159,7 +154,8 @@ fn compute_method_recursive_cost(
         sum += compute_method_recursive_cost(child, methods);
     }
 
-    if !node.children.is_empty() {
+    // Only include entries that look like function names (contain "::")
+    if !node.children.is_empty() && is_function_name(&node.text) {
         insert_or_add(methods, node.text.clone(), sum);
     }
     sum
@@ -177,7 +173,8 @@ fn compute_method_self_cost(
         sum += child.val;
     }
 
-    if !node.children.is_empty() {
+    // Only include entries that look like function names (contain "::")
+    if !node.children.is_empty() && is_function_name(&node.text) {
         insert_or_add(methods, node.text.clone(), sum);
     }
 }

aptos-move/aptos-gas-profiling/src/erased.rs

@@ -133,29 +133,25 @@ impl ExecutionGasEvent {
         use ExecutionGasEvent::*;
 
         match self {
-            Loc(offset) => Node::new(format!("@{}", offset), 0),
-            Bytecode { op, cost } => Node::new(format!("{:?}", op).to_ascii_lowercase(), *cost),
+            Loc(offset) => Node::new(format!("@{offset}"), 0),
+            Bytecode { op, cost } => Node::new(format!("{op:?}").to_ascii_lowercase(), *cost),
             Call(frame) => frame.to_erased(keep_generic_types),
             CallNative {
                 module_id,
                 fn_name,
                 ty_args,
                 cost,
-            } => Node::new(
-                format!(
-                    "{}",
-                    Render(&(
-                        module_id,
-                        fn_name.as_ident_str(),
-                        if keep_generic_types {
-                            ty_args.as_slice()
-                        } else {
-                            &[]
-                        }
-                    ))
-                ),
-                *cost,
-            ),
+            } => {
+                let ty_args = if keep_generic_types {
+                    ty_args.as_slice()
+                } else {
+                    &[]
+                };
+                Node::new(
+                    Render(&(module_id, fn_name.as_ident_str(), ty_args)).to_string(),
+                    *cost,
+                )
+            },
             LoadResource { addr, ty, cost } => Node::new(
                 format!("load<{}::{}>", Render(addr), ty.to_canonical_string()),
                 *cost,
@@ -175,26 +171,19 @@ impl CallFrame {
                 name,
                 ty_args,
             } => {
-                format!(
-                    "{}",
-                    Render(&(
-                        module_id,
-                        name.as_ident_str(),
-                        if keep_generic_types {
-                            ty_args.as_slice()
-                        } else {
-                            &[]
-                        }
-                    ))
-                )
+                let ty_args = if keep_generic_types {
+                    ty_args.as_slice()
+                } else {
+                    &[]
+                };
+                Render(&(module_id, name.as_ident_str(), ty_args)).to_string()
             },
         };
 
         let children = self
             .events
             .iter()
-            .map(|event| event.to_erased(keep_generic_types))
-            .collect::<Vec<_>>();
+            .map(|event| event.to_erased(keep_generic_types));
 
         Node::new_with_children(name, 0, children)
     }
@@ -247,7 +236,7 @@ impl ExecutionAndIOCosts {
 
         TypeErasedExecutionAndIoCosts {
             gas_scaling_factor: self.gas_scaling_factor,
-            total: self.total,
+            total: self.total(),
             tree: Node::new_with_children("execution & IO (gas unit, full trace)", 0, nodes),
         }
     }

aptos-move/aptos-gas-profiling/src/flamegraph.rs

@@ -5,10 +5,46 @@ use crate::{
     log::{CallFrame, ExecutionAndIOCosts, ExecutionGasEvent, StorageFees},
     render::Render,
 };
-use inferno::flamegraph::TextTruncateDirection;
+use inferno::flamegraph::{Options, TextTruncateDirection};
 use move_core_types::gas_algebra::InternalGas;
 use regex::Captures;
 
+/// Creates default flamegraph options with the given title.
+fn flamegraph_options(title: String) -> Options<'static> {
+    let mut options = Options::default();
+    options.flame_chart = true;
+    options.text_truncate_direction = TextTruncateDirection::Right;
+    options.color_diffusion = true;
+    options.title = title;
+    options
+}
+
+/// Generates a flamegraph SVG from folded stack lines, applying a regex replacement to the output.
+fn generate_flamegraph(
+    lines: Vec<String>,
+    title: String,
+    replace_samples: impl Fn(&Captures) -> String,
+) -> anyhow::Result<Option<Vec<u8>>> {
+    if lines.is_empty() {
+        return Ok(None);
+    }
+
+    let mut options = flamegraph_options(title);
+    let mut graph_content = vec![];
+    inferno::flamegraph::from_lines(
+        &mut options,
+        lines.iter().rev().map(|s| s.as_str()),
+        &mut graph_content,
+    )?;
+
+    let graph_content = String::from_utf8_lossy(&graph_content);
+    let re = regex::Regex::new("([1-9][0-9]*(,[0-9]+)*) samples")
+        .expect("should be able to build regex successfully");
+    let graph_content = re.replace_all(&graph_content, replace_samples);
+
+    Ok(Some(graph_content.as_bytes().to_vec()))
+}
+
 #[derive(Debug)]
 struct LineBuffer(Vec<String>);
 
@@ -56,43 +92,16 @@ impl StorageFees {
         lines.into_inner()
     }
 
-    /// Tries to generate a flamegraph from the execution log.
+    /// Tries to generate a flamegraph from the storage fee log.
     /// None will be returned if the log is empty.
     pub fn to_flamegraph(&self, title: String) -> anyhow::Result<Option<Vec<u8>>> {
-        let lines = self.to_folded_stack_lines();
-
-        if lines.is_empty() {
-            return Ok(None);
-        }
-
-        let mut options = inferno::flamegraph::Options::default();
-        options.flame_chart = true;
-        options.text_truncate_direction = TextTruncateDirection::Right;
-        options.color_diffusion = true;
-        options.title = title;
-
-        let mut graph_content = vec![];
-        inferno::flamegraph::from_lines(
-            &mut options,
-            lines.iter().rev().map(|s| s.as_str()),
-            &mut graph_content,
-        )?;
-        let graph_content = String::from_utf8_lossy(&graph_content);
-
-        // Inferno does not allow us to customize some of the text in the resulting graph,
-        // so we have to do it through regex replacement.
-        let re = regex::Regex::new("([1-9][0-9]*(,[0-9]+)*) samples")
-            .expect("should be able to build regex successfully");
-        let graph_content = re.replace_all(&graph_content, |caps: &Captures| {
+        generate_flamegraph(self.to_folded_stack_lines(), title, |caps| {
             let count: u64 = caps[1]
                 .replace(',', "")
                 .parse()
                 .expect("should be able parse count as u64");
-
             format!("{} Octa", count)
-        });
-
-        Ok(Some(graph_content.as_bytes().to_vec()))
+        })
     }
 }
 
@@ -201,47 +210,20 @@ impl ExecutionAndIOCosts {
     /// Tries to generate a flamegraph from the execution log.
     /// None will be returned if the log is empty.
     pub fn to_flamegraph(&self, title: String) -> anyhow::Result<Option<Vec<u8>>> {
-        let lines = self.to_folded_stack_lines();
-
-        if lines.is_empty() {
-            return Ok(None);
-        }
-
-        let mut options = inferno::flamegraph::Options::default();
-        options.flame_chart = true;
-        options.text_truncate_direction = TextTruncateDirection::Right;
-        options.color_diffusion = true;
-        options.title = title;
-
-        let mut graph_content = vec![];
-        inferno::flamegraph::from_lines(
-            &mut options,
-            lines.iter().rev().map(|s| s.as_str()),
-            &mut graph_content,
-        )?;
-        let graph_content = String::from_utf8_lossy(&graph_content);
-
-        // Inferno does not allow us to customize some of the text in the resulting graph,
-        // so we have to do it through regex replacement.
-        let re = regex::Regex::new("([1-9][0-9]*(,[0-9]+)*) samples")
-            .expect("should be able to build regex successfully");
-        let graph_content = re.replace_all(&graph_content, |caps: &Captures| {
+        let scaling_factor = u64::from(self.gas_scaling_factor) as f64;
+        generate_flamegraph(self.to_folded_stack_lines(), title, |caps| {
             let count: u64 = caps[1]
                 .replace(',', "")
                 .parse()
                 .expect("should be able parse count as u64");
-
-            let count_scaled = count as f64 / u64::from(self.gas_scaling_factor) as f64;
-
+            let count_scaled = count as f64 / scaling_factor;
             format!(
                 "{} gas units",
                 crate::misc::strip_trailing_zeros_and_decimal_point(&format!(
                     "{:.8}",
                     count_scaled
                 ))
             )
-        });
-
-        Ok(Some(graph_content.as_bytes().to_vec()))
+        })
     }
 }

aptos-move/aptos-gas-profiling/src/log.rs

@@ -129,7 +129,10 @@ impl Dependency {
 #[derive(Debug, Clone)]
 pub struct ExecutionAndIOCosts {
     pub gas_scaling_factor: GasScalingFactor,
-    pub total: InternalGas,
+    /// Execution gas - corresponds to FeeStatement::execution_gas_units.
+    pub execution_gas: InternalGas,
+    /// IO gas - corresponds to FeeStatement::io_gas_units.
+    pub io_gas: InternalGas,
 
     pub intrinsic_cost: InternalGas,
     pub keyless_cost: InternalGas,
@@ -253,6 +256,11 @@ impl StorageFees {
 }
 
 impl ExecutionAndIOCosts {
+    /// Returns the total execution + IO gas.
+    pub fn total(&self) -> InternalGas {
+        self.execution_gas + self.io_gas
+    }
+
     #[allow(clippy::needless_lifetimes)]
     pub fn gas_events<'a>(&'a self) -> GasEventIter<'a> {
         GasEventIter {
@@ -295,10 +303,10 @@ impl ExecutionAndIOCosts {
             total += write.cost;
         }
 
-        if total != self.total {
+        if total != self.total() {
             panic!(
                 "Execution & IO costs do not add up. Check if the gas meter & the gas profiler have been implemented correctly. From gas meter: {}. Calculated: {}.",
-                self.total, total
+                self.total(), total
             );
         }
     }

aptos-move/aptos-gas-profiling/src/profiler.rs

@@ -718,8 +718,11 @@ where
     fn charge_slh_dsa_sha2_128s(&mut self) -> VMResult<()> {
         let (cost, res) = self.delegate_charge(|base| base.charge_slh_dsa_sha2_128s());
 
-        self.slh_dsa_sha2_128s_cost =
-            Some(self.slh_dsa_sha2_128s_cost.unwrap_or_else(|| 0.into()) + cost);
+        self.slh_dsa_sha2_128s_cost = Some(
+            self.slh_dsa_sha2_128s_cost
+                .unwrap_or_else(InternalGas::zero)
+                + cost,
+        );
 
         res
     }
@@ -738,10 +741,13 @@ where
 
         let exec_io = ExecutionAndIOCosts {
             gas_scaling_factor: self.base.gas_unit_scaling_factor(),
-            total: self.algebra().execution_gas_used() + self.algebra().io_gas_used(),
-            intrinsic_cost: self.intrinsic_cost.unwrap_or_else(|| 0.into()),
-            keyless_cost: self.keyless_cost.unwrap_or_else(|| 0.into()),
-            slh_dsa_sha2_128s_cost: self.slh_dsa_sha2_128s_cost.unwrap_or_else(|| 0.into()),
+            execution_gas: self.algebra().execution_gas_used(),
+            io_gas: self.algebra().io_gas_used(),
+            intrinsic_cost: self.intrinsic_cost.unwrap_or_else(InternalGas::zero),
+            keyless_cost: self.keyless_cost.unwrap_or_else(InternalGas::zero),
+            slh_dsa_sha2_128s_cost: self
+                .slh_dsa_sha2_128s_cost
+                .unwrap_or_else(InternalGas::zero),
             dependencies: self.dependencies,
             call_graph: self.frames.pop().expect("frame must exist"),
             transaction_transient: self.transaction_transient,
@@ -751,12 +757,12 @@ where
         exec_io.assert_consistency();
 
         let storage = self.storage_fees.unwrap_or_else(|| StorageFees {
-            total: 0.into(),
-            total_refund: 0.into(),
+            total: Fee::zero(),
+            total_refund: Fee::zero(),
             write_set_storage: vec![],
             events: vec![],
-            event_discount: 0.into(),
-            txn_storage: 0.into(),
+            event_discount: Fee::zero(),
+            txn_storage: Fee::zero(),
         });
         storage.assert_consistency();