Trie stats, query params, and better usage docs

Cargo.toml

@@ -1,11 +1,11 @@
 [workspace]
-members = [ "common", "trie_diff", "trie_query"]
+members = [ "common", "trie_diff", "trie_query", "trie_stats" ]
 resolver = "2"
 
 [workspace.dependencies]
 anyhow = { version = "1.0.75", features = ["backtrace"] }
 clap = { version = "4.4.6", features = ["derive", "env"] }
-mpt_trie = { version = "0.1.0", git = "https://github.com/0xPolygonZero/zk_evm.git", branch = "robins_not_so_awesome_branch_for_brendan" }
+mpt_trie = "0.3.0"
 pretty_env_logger = "0.5.0"
-trace_decoder = { version = "0.1.0", git = "https://github.com/0xPolygonZero/zk_evm.git", branch = "robins_not_so_awesome_branch_for_brendan" }
+trace_decoder = "0.4.0"
 serde_json = "1.0.114"

common/Cargo.toml

@@ -6,6 +6,7 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+clap = { workspace = true }
 mpt_trie = { workspace = true }
 hex = "0.4.3"
 

common/src/lib.rs

@@ -6,15 +6,22 @@ use trace_decoder::{
     trace_protocol::TrieCompact,
 };
 
+pub const TRIE_PATH_DOC_STR: &str = r#"
+Each path must point to a file that is either:
+- (`*.compact`) Compact encoding (hex string) of an MPT trie (spec: https://gist.github.com/mandrigin/ff7eccf30d0ef9c572bafcb0ab665cff).
+- (`*.json`) A serialized `HashedPartialTrie`.
+"#;
+
 pub fn read_input_from_file(p: &Path) -> HashedPartialTrie {
     match p
         .extension()
         .map(|s| s.to_str().unwrap_or_default())
         .unwrap_or_default()
     {
+        // For now, if we ever get compact, we're going to just use the state trie.
         "compact" => {
             let out = process_compact_prestate_debug(read_compact_from_file(p)).unwrap();
-            out.witness_out.tries.state
+            out.witness_out.state_trie
         }
         "json" => read_json_trie_from_file(p),
         _ => panic!(

trie_diff/src/main.rs

@@ -1,11 +1,15 @@
 use std::path::PathBuf;
 
 use clap::Parser;
-use common::read_input_from_file;
+use common::{read_input_from_file, TRIE_PATH_DOC_STR};
 use mpt_trie::debug_tools::diff::create_diff_between_tries;
 use mpt_trie::partial_trie::PartialTrie;
 
 #[derive(Debug, Parser)]
+#[command(after_help = TRIE_PATH_DOC_STR)]
+/// Attempts to find the lowest possible point where two given tries differ.
+///
+/// Takes in two paths to tries that we want to include in the diff.
 struct ProgArgs {
     a_path: PathBuf,
     b_path: PathBuf,

trie_query/src/main.rs

@@ -1,16 +1,51 @@
 use std::path::PathBuf;
 
 use clap::Parser;
-use common::read_input_from_file;
+use common::{read_input_from_file, TRIE_PATH_DOC_STR};
 use mpt_trie::{
     debug_tools::query::{get_path_from_query, DebugQueryParamsBuilder},
     nibbles::Nibbles,
 };
 
+const QUERY_CONFIG_STR: &str = "Query Config";
+
+/// Displays a "trace" of a key query for a given key, where the trace contains the nodes encountered.
 #[derive(Debug, Parser)]
+#[command(after_help = TRIE_PATH_DOC_STR)]
 struct ProgArgs {
+    /// Path to the trie to query against.
     trie_path: PathBuf,
+
+    /// The key to query in hex format (with or without the "0x").
     key: Nibbles,
+
+    /// Config for the query output.
+    #[command(flatten)]
+    config: QueryConfig,
+}
+
+#[derive(Parser, Debug)]
+struct QueryConfig {
+    /// Include the key piece (if the node contains a piece of the key) contained in each node.
+    #[arg(short = 'k', long, default_value_t = true, help_heading = QUERY_CONFIG_STR)]
+    include_key_piece_per_node: bool,
+
+    /// Include the node types in the trace.
+    #[arg(short = 't', long, default_value_t = true, help_heading = QUERY_CONFIG_STR)]
+    include_node_type: bool,
+
+    /// Include additional info that is specific to a given node type.
+    #[arg(short = 's', long, default_value_t = false, help_heading = QUERY_CONFIG_STR)]
+    include_node_specific_values: bool,
+}
+
+impl From<QueryConfig> for DebugQueryParamsBuilder {
+    fn from(v: QueryConfig) -> Self {
+        DebugQueryParamsBuilder::default()
+            .print_key_pieces(v.include_key_piece_per_node)
+            .print_node_type(v.include_node_type)
+            .print_node_specific_values(v.include_node_specific_values)
+    }
 }
 
 fn main() -> anyhow::Result<()> {
@@ -19,9 +54,7 @@ fn main() -> anyhow::Result<()> {
     let p_args = ProgArgs::parse();
     let trie = read_input_from_file(&p_args.trie_path);
 
-    let query = DebugQueryParamsBuilder::default()
-        .print_node_specific_values(true)
-        .build(p_args.key);
+    let query = DebugQueryParamsBuilder::from(p_args.config).build(p_args.key);
     let res = get_path_from_query(&trie, query);
 
     println!("{}", res);

trie_stats/Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "trie_stats"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+common = { path = "../common" }
+
+anyhow = { workspace = true }
+clap = { workspace = true }
+mpt_trie ={ workspace = true }
+pretty_env_logger = { workspace = true }

trie_stats/src/main.rs

@@ -0,0 +1,47 @@
+use std::path::PathBuf;
+
+use clap::Parser;
+use common::{read_input_from_file, TRIE_PATH_DOC_STR};
+use mpt_trie::debug_tools::stats::get_trie_stats_with_name;
+
+/// Analyze one or two tries and output stats.
+#[derive(Debug, Parser)]
+#[command(after_help = TRIE_PATH_DOC_STR)]
+struct ProgArgs {
+    /// The primary trie to compare against.
+    trie_path: PathBuf,
+
+    /// If a second trie is provided, a comparison will be made between the two tries along with printout out their individual stats.
+    other_trie_path: Option<PathBuf>,
+}
+
+fn main() {
+    pretty_env_logger::init();
+
+    let p_args = ProgArgs::parse();
+    let trie = read_input_from_file(&p_args.trie_path);
+    let trie_name = p_args
+        .trie_path
+        .file_name()
+        .map(|os_str| os_str.to_string_lossy())
+        .unwrap_or("Trie".into());
+
+    let stats = get_trie_stats_with_name(&trie, trie_name.to_string());
+    println!("{}", stats);
+
+    if let Some(other_trie_path) = p_args.other_trie_path {
+        // A second trie was passed in. Print its stats and also do a comparison.
+        let other_trie = read_input_from_file(&other_trie_path);
+        let other_trie_name = other_trie_path
+            .file_name()
+            .map(|os_str| os_str.to_string_lossy())
+            .unwrap_or("Trie".into());
+        let other_trie_stats = get_trie_stats_with_name(&other_trie, other_trie_name.into());
+
+        println!("{}", other_trie_stats);
+
+        let comparison = stats.compare(&other_trie_stats);
+
+        println!("{}", comparison);
+    }
+}