Add additional mechanisms to introspect state witnesses

core/primitives/src/challenge.rs

@@ -31,7 +31,8 @@ impl Debug for PartialState {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
             PartialState::TrieValues(values) => {
-                f.write_str(&format!("{} trie values", values.len()))
+                let size = bytesize::ByteSize(values.iter().map(|v| v.len() as u64).sum());
+                f.write_str(&format!("{} trie values sized {}", values.len(), size))
             }
         }
     }

core/primitives/src/stateless_validation/state_witness.rs

@@ -145,7 +145,7 @@ pub struct ChunkStateWitness {
     pub new_transactions: Vec<SignedTransaction>,
     pub new_transactions_validation_state: PartialState,
     // TODO(stateless_validation): Deprecate once we send state witness in parts.
-    signature_differentiator: SignatureDifferentiator,
+    pub signature_differentiator: SignatureDifferentiator,
 }
 
 impl ChunkStateWitness {

tools/state-viewer/src/latest_witnesses.rs

@@ -1,14 +1,18 @@
+use std::collections::BTreeMap;
 use std::path::{Path, PathBuf};
 use std::rc::Rc;
 
+use borsh::BorshDeserialize as _;
 use near_chain::runtime::NightshadeRuntime;
 use near_chain::stateless_validation::processing_tracker::ProcessingDoneTracker;
 use near_chain::{Chain, ChainGenesis, ChainStore, DoomslugThresholdMode};
 use near_epoch_manager::shard_tracker::{ShardTracker, TrackedConfig};
 use near_epoch_manager::EpochManager;
+use near_primitives::challenge::PartialState;
+use near_primitives::hash::{hash, CryptoHash};
 use near_primitives::stateless_validation::state_witness::ChunkStateWitness;
 use near_primitives::types::EpochId;
-use near_store::Store;
+use near_store::{RawTrieNodeWithSize, Store};
 use near_time::Clock;
 use nearcore::NearConfig;
 use nearcore::NightshadeRuntimeExt;
@@ -48,13 +52,76 @@ pub struct DumpWitnessesCmd {
 
 #[derive(clap::Subcommand)]
 enum DumpWitnessesMode {
-    /// Pretty-print on screen using the "{:#?}" formatting.
-    Pretty,
+    /// Pretty-print the contents on screen.
+    Pretty {
+        #[clap(long)]
+        decode_tries: bool,
+    },
     /// Saves the raw &[u8] of each witness to the given directory.
     Binary { output_dir: PathBuf },
 }
 
 impl DumpWitnessesCmd {
+    fn pretty_print_partial_state(&self, partial_state: &PartialState) {
+        let PartialState::TrieValues(nodes) = partial_state;
+        let mut nmap = nodes.iter().map(|n| (hash(&n), (None, &**n))).collect::<BTreeMap<_, _>>();
+        fn compute_size(nmap: &mut BTreeMap<CryptoHash, (Option<u64>, &[u8])>, h: &CryptoHash) {
+            let Some((sz, val)) = nmap.get(h) else { return };
+            if sz.is_some() {
+                return;
+            }
+            let self_size = val.len() as u64;
+            let Ok(node) = RawTrieNodeWithSize::try_from_slice(val) else {
+                nmap.get_mut(h).unwrap().0.replace(self_size);
+                return;
+            };
+            let size = match node.node {
+                near_store::RawTrieNode::Leaf(_, r) => r.length as u64,
+                near_store::RawTrieNode::BranchNoValue(ch) => ch
+                    .iter()
+                    .map(|(_, c)| {
+                        compute_size(nmap, c);
+                        nmap.get(c).map(|v| v.0.unwrap()).unwrap_or(0)
+                    })
+                    .sum(),
+                near_store::RawTrieNode::BranchWithValue(v, ch) => {
+                    ch.iter()
+                        .map(|(_, c)| {
+                            compute_size(nmap, c);
+                            nmap.get(c).map(|v| v.0.unwrap()).unwrap_or(0)
+                        })
+                        .sum::<u64>()
+                        + v.length as u64
+                }
+                near_store::RawTrieNode::Extension(_, c) => {
+                    compute_size(nmap, &c);
+                    nmap.get(&c).map(|v| v.0.unwrap()).unwrap_or(0)
+                }
+            };
+            nmap.get_mut(h).unwrap().0.replace(self_size + size);
+        }
+
+        for n in nmap.keys().copied().collect::<Vec<_>>() {
+            compute_size(&mut nmap, &n);
+        }
+
+        for node in nodes {
+            let hash = hash(node);
+            let Ok(node) = RawTrieNodeWithSize::try_from_slice(node) else {
+                println!("  {}: Not RawTrieNodeWithSize (size: {})", hash, node.len());
+                continue;
+            };
+
+            println!(
+                "  {}: .memory_usage: {}, .actual_cumulative_memory: {}, .node: {:?}",
+                hash,
+                bytesize::ByteSize(node.memory_usage),
+                bytesize::ByteSize(nmap.get(&hash).unwrap().0.unwrap()),
+                node.node,
+            );
+        }
+    }
+
     pub(crate) fn run(&self, near_config: NearConfig, store: Store) {
         let chain_store =
             Rc::new(ChainStore::new(store, near_config.genesis.config.genesis_height, false));
@@ -77,9 +144,55 @@ impl DumpWitnessesCmd {
                 witness.epoch_id
             );
             match self.mode {
-                DumpWitnessesMode::Pretty => {
-                    println!("{:#?}", witness);
-                    println!("");
+                DumpWitnessesMode::Pretty { decode_tries } => {
+                    let ChunkStateWitness {
+                        chunk_producer,
+                        epoch_id,
+                        chunk_header,
+                        main_state_transition,
+                        source_receipt_proofs,
+                        applied_receipts_hash,
+                        transactions,
+                        implicit_transitions,
+                        new_transactions,
+                        new_transactions_validation_state,
+                        signature_differentiator,
+                    } = witness;
+                    println!("chunk_producer: {:#?}", chunk_producer);
+                    println!("epoch_id: {:#?}", epoch_id);
+                    println!("chunk_header: {:#?}", chunk_header);
+                    println!(
+                        "main_state_transition.block_hash: {:#?}",
+                        main_state_transition.block_hash
+                    );
+                    println!(
+                        "main_state_transition.post_state_root: {:#?}",
+                        main_state_transition.post_state_root
+                    );
+                    if decode_tries {
+                        println!("main_state_transition.base_state: (continued)",);
+                        self.pretty_print_partial_state(&main_state_transition.base_state);
+                    } else {
+                        println!(
+                            "main_state_transition.base_state: {:#?}",
+                            main_state_transition.base_state
+                        );
+                    }
+                    println!("source_receipt_proofs: {:#?}", source_receipt_proofs);
+                    println!("applied_receipts_hash: {:#?}", applied_receipts_hash);
+                    println!("transactions: {:#?}", transactions);
+                    println!("implicit_transitions: {:#?}", implicit_transitions);
+                    println!("new_transactions: {:#?}", new_transactions);
+                    if decode_tries {
+                        println!("new_transactions_validation_state: (continued)",);
+                        self.pretty_print_partial_state(new_transactions_validation_state);
+                    } else {
+                        println!(
+                            "new_transactions_validation_state: {:#?}",
+                            new_transactions_validation_state
+                        );
+                    }
+                    println!("signature_differentiator: {:#?}", signature_differentiator);
                 }
                 DumpWitnessesMode::Binary { ref output_dir } => {
                     let file_name = format!(