Hide address and selector (#41)

Co-authored-by: Piotr Magiera <56825108+piotmag769@users.noreply.github.com>

Cargo.toml

@@ -22,3 +22,4 @@ assert_fs = "1.1.1"
 project-root = "0.2.2"
 snapbox = "0.4.16"
 prost-build = { version = "0.12" }
+test-case = "3.3.1"

crates/cairo-profiler/Cargo.toml

@@ -23,6 +23,7 @@ trace-data = { path="../trace-data" }
 assert_fs.workspace = true
 project-root.workspace = true
 snapbox.workspace = true
+test-case.workspace = true
 
 [build-dependencies]
 prost-build.workspace = true

crates/cairo-profiler/src/main.rs

@@ -26,6 +26,10 @@ struct Cli {
     /// Path to the output file
     #[arg(short, long, default_value = "profile.pb.gz")]
     output_path: Utf8PathBuf,
+
+    /// Show contract addresses and function selectors in a trace tree
+    #[arg(long)]
+    show_details: bool,
 }
 
 fn main() -> Result<()> {
@@ -35,7 +39,7 @@ fn main() -> Result<()> {
         .context("Failed to read call trace from a file")?;
     let serialized_trace: CallTrace =
         serde_json::from_str(&data).context("Failed to deserialize call trace")?;
-    let samples = collect_samples_from_trace(&serialized_trace);
+    let samples = collect_samples_from_trace(&serialized_trace, cli.show_details);
     let resources_keys = collect_resources_keys(&samples);
 
     let profile = build_profile(&samples, &resources_keys);

crates/cairo-profiler/src/trace_reader.rs

@@ -85,23 +85,40 @@ impl Sample {
     }
 }
 
+/// `contract_name` and `function_name` are always present (in case they are not in trace we just
+/// set `<unknown>` string)
+/// `address` and `selector` are optional and set if `--show-details` flag is enabled
+/// or names are unknown
 pub struct EntryPointId {
-    address: String,
-    selector: String,
-    contract_name: Option<String>,
-    function_name: Option<String>,
+    address: Option<String>,
+    selector: Option<String>,
+    contract_name: String,
+    function_name: String,
 }
 
 impl EntryPointId {
     fn from(
         contract_name: Option<String>,
         function_name: Option<String>,
         contract_address: ContractAddress,
-        selector: EntryPointSelector,
+        function_selector: EntryPointSelector,
+        show_details: bool,
     ) -> Self {
+        let (contract_name, address) = match contract_name {
+            Some(name) if show_details => (name, Some(contract_address.0)),
+            Some(name) => (name, None),
+            None => (String::from("<unknown>"), Some(contract_address.0)),
+        };
+
+        let (function_name, selector) = match function_name {
+            Some(name) if show_details => (name, Some(function_selector.0)),
+            Some(name) => (name, None),
+            None => (String::from("<unknown>"), Some(function_selector.0)),
+        };
+
         EntryPointId {
-            address: contract_address.0,
-            selector: selector.0,
+            address,
+            selector,
             contract_name,
             function_name,
         }
@@ -110,26 +127,27 @@ impl EntryPointId {
 
 impl Display for EntryPointId {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        let contract_name = self
-            .contract_name
-            .clone()
-            .unwrap_or(String::from("<unknown>"));
-        let function_name = self
-            .function_name
-            .clone()
-            .unwrap_or(String::from("<unknown>"));
+        let contract_address = match &self.address {
+            None => String::new(),
+            Some(address) => format!("Address: {address}\n"),
+        };
+        let selector = match &self.selector {
+            None => String::new(),
+            Some(selector) => format!("Selector: {selector}\n"),
+        };
+
         write!(
             f,
-            "Contract address: {}\n Selector: {}\nContract name: {}\nFunction name: {}\n",
-            self.address, self.selector, contract_name, function_name
+            "Contract: {}\n{contract_address}Function: {}\n{selector}",
+            self.contract_name, self.function_name
         )
     }
 }
 
-pub fn collect_samples_from_trace(trace: &CallTrace) -> Vec<Sample> {
+pub fn collect_samples_from_trace(trace: &CallTrace, show_details: bool) -> Vec<Sample> {
     let mut samples = vec![];
     let mut current_path = vec![];
-    collect_samples(&mut samples, &mut current_path, trace);
+    collect_samples(&mut samples, &mut current_path, trace, show_details);
     samples
 }
 
@@ -204,17 +222,20 @@ fn collect_samples<'a>(
     samples: &mut Vec<Sample>,
     current_call_stack: &mut Vec<EntryPointId>,
     trace: &'a CallTrace,
+    show_details: bool,
 ) -> &'a ExecutionResources {
     current_call_stack.push(EntryPointId::from(
         trace.entry_point.contract_name.clone(),
         trace.entry_point.function_name.clone(),
         trace.entry_point.contract_address.clone(),
         trace.entry_point.entry_point_selector.clone(),
+        show_details,
     ));
 
     let mut children_resources = ExecutionResources::default();
     for sub_trace in &trace.nested_calls {
-        children_resources += &collect_samples(samples, current_call_stack, sub_trace);
+        children_resources +=
+            &collect_samples(samples, current_call_stack, sub_trace, show_details);
     }
 
     samples.push(Sample {

crates/cairo-profiler/tests/e2e.rs

@@ -1,8 +1,9 @@
 use assert_fs::fixture::PathCopy;
 use snapbox::cmd::{cargo_bin, Command as SnapboxCommand};
+use test_case::test_case;
 
 #[test]
-fn simple_package() {
+fn output_path() {
     let project_root = project_root::get_project_root().unwrap();
     let temp_dir = assert_fs::TempDir::new().unwrap();
     temp_dir
@@ -15,16 +16,16 @@ fn simple_package() {
     SnapboxCommand::new(cargo_bin!("cairo-profiler"))
         .current_dir(&temp_dir)
         .arg("./trace.json")
+        .args(["-o", "my/output/dir/my_file.pb.gz"])
         .assert()
         .success();
 
-    assert!(temp_dir.join("profile.pb.gz").exists());
-
-    // TODO run pprof here
+    assert!(temp_dir.join("my/output/dir/my_file.pb.gz").exists());
 }
 
-#[test]
-fn output_path() {
+#[test_case(&["trace.json", "--show-details"]; "with details")]
+#[test_case(&["trace.json"]; "without details")]
+fn simple_package(args: &[&str]) {
     let project_root = project_root::get_project_root().unwrap();
     let temp_dir = assert_fs::TempDir::new().unwrap();
     temp_dir
@@ -36,10 +37,11 @@ fn output_path() {
 
     SnapboxCommand::new(cargo_bin!("cairo-profiler"))
         .current_dir(&temp_dir)
-        .arg("./trace.json")
-        .args(["-o", "my/output/dir/my_file.pb.gz"])
+        .args(args)
         .assert()
         .success();
 
-    assert!(temp_dir.join("my/output/dir/my_file.pb.gz").exists());
+    assert!(temp_dir.join("profile.pb.gz").exists());
+
+    // TODO run pprof here
 }