From 7c60ae70fe8bedacb78f8ff5783b9517c94365c8 Mon Sep 17 00:00:00 2001
From: Alexey Shekhirin <5773434+shekhirin@users.noreply.github.com>
Date: Thu, 30 Jan 2025 21:33:38 +0000
Subject: [PATCH] better metrics

---
 crates/engine/tree/src/tree/root.rs | 25 +++++++++++++++----------
 crates/trie/common/src/proofs.rs    |  4 ++++
 crates/trie/parallel/src/proof.rs   |  4 ++--
 3 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/crates/engine/tree/src/tree/root.rs b/crates/engine/tree/src/tree/root.rs
index 6f765da4924b..bf839c0901d5 100644
--- a/crates/engine/tree/src/tree/root.rs
+++ b/crates/engine/tree/src/tree/root.rs
@@ -440,10 +440,10 @@ where
 struct StateRootTaskMetrics {
     /// Histogram of proof calculation durations.
     pub proof_calculation_duration_histogram: Histogram,
-    /// Histogram of proof calculation account targets.
-    pub proof_calculation_account_targets_histogram: Histogram,
-    /// Histogram of proof calculation storage targets.
-    pub proof_calculation_storage_targets_histogram: Histogram,
+    /// Histogram of proof calculation targets with a ccounts.
+    pub proof_calculation_with_account_targets_histogram: Histogram,
+    /// Histogram of proof calculation targets with only storage slots.
+    pub proof_calculation_only_storage_targets_histogram: Histogram,
 
     /// Histogram of sparse trie update durations.
     pub sparse_trie_update_duration_histogram: Histogram,
@@ -689,16 +689,21 @@ where
                         self.metrics
                             .proof_calculation_duration_histogram
                             .record(proof_calculated.elapsed);
-                        self.metrics
-                            .proof_calculation_account_targets_histogram
-                            .record(proof_calculated.update.targets.len() as f64);
-                        self.metrics.proof_calculation_storage_targets_histogram.record(
+                        self.metrics.proof_calculation_with_account_targets_histogram.record(
+                            proof_calculated
+                                .update
+                                .targets
+                                .values()
+                                .filter(|slots| slots.is_with_account())
+                                .count() as f64,
+                        );
+                        self.metrics.proof_calculation_only_storage_targets_histogram.record(
                             proof_calculated
                                 .update
                                 .targets
                                 .values()
-                                .map(|targets| targets.len() as f64)
-                                .sum::<f64>(),
+                                .filter(|slots| slots.is_only_storage())
+                                .count() as f64,
                         );
 
                         debug!(
diff --git a/crates/trie/common/src/proofs.rs b/crates/trie/common/src/proofs.rs
index dd297ce789d9..ee3cca72512a 100644
--- a/crates/trie/common/src/proofs.rs
+++ b/crates/trie/common/src/proofs.rs
@@ -42,6 +42,10 @@ impl MultiProofAccountTarget {
         matches!(self, Self::WithAccount(_))
     }
 
+    pub const fn is_only_storage(&self) -> bool {
+        matches!(self, Self::OnlyStorage(_))
+    }
+
     pub fn is_empty(&self) -> bool {
         match self {
             Self::WithAccount(set) | Self::OnlyStorage(set) => set.is_empty(),
diff --git a/crates/trie/parallel/src/proof.rs b/crates/trie/parallel/src/proof.rs
index 913f3ff401ca..637f6efbc7f6 100644
--- a/crates/trie/parallel/src/proof.rs
+++ b/crates/trie/parallel/src/proof.rs
@@ -96,8 +96,8 @@ where
             B256HashMap<B256HashSet>,
             B256HashMap<B256HashSet>,
         ) = targets.into_iter().partition_map(|(key, target)| match target {
-            MultiProofAccountTarget::WithAccount(hash_set) => Either::Left((key, hash_set)),
-            MultiProofAccountTarget::OnlyStorage(hash_set) => Either::Right((key, hash_set)),
+            MultiProofAccountTarget::WithAccount(slots) => Either::Left((key, slots)),
+            MultiProofAccountTarget::OnlyStorage(slots) => Either::Right((key, slots)),
         });
 
         // Extend prefix sets with targets