From e2015143b35d9d6c2cf180bec752673767063f0f Mon Sep 17 00:00:00 2001
From: Roman Krasiuk <rokrassyuk@gmail.com>
Date: Tue, 25 Jun 2024 04:57:28 -0700
Subject: [PATCH] chore(trie): add helpers to return trie keys as variants
 (#9075)

---
 crates/trie/trie/src/updates.rs | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/crates/trie/trie/src/updates.rs b/crates/trie/trie/src/updates.rs
index afad628cd7..53830fd8a3 100644
--- a/crates/trie/trie/src/updates.rs
+++ b/crates/trie/trie/src/updates.rs
@@ -23,6 +23,35 @@ pub enum TrieKey {
     StorageTrie(B256),
 }
 
+impl TrieKey {
+    /// Returns reference to account node key if the key is for [`Self::AccountNode`].
+    pub const fn as_account_node_key(&self) -> Option<&StoredNibbles> {
+        if let Self::AccountNode(nibbles) = &self {
+            Some(nibbles)
+        } else {
+            None
+        }
+    }
+
+    /// Returns reference to storage node key if the key is for [`Self::StorageNode`].
+    pub const fn as_storage_node_key(&self) -> Option<(&B256, &StoredNibblesSubKey)> {
+        if let Self::StorageNode(key, subkey) = &self {
+            Some((key, subkey))
+        } else {
+            None
+        }
+    }
+
+    /// Returns reference to storage trie key if the key is for [`Self::StorageTrie`].
+    pub const fn as_storage_trie_key(&self) -> Option<&B256> {
+        if let Self::StorageTrie(key) = &self {
+            Some(key)
+        } else {
+            None
+        }
+    }
+}
+
 /// The operation to perform on the trie.
 #[derive(PartialEq, Eq, Debug, Clone)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]