POC of trie keys to abstract hashing and prefixes

crates/storage/layering.rs

@@ -284,7 +284,7 @@ impl TrieWrapper {
         db: Box<dyn TrieDB>,
         prefix: Option<H256>,
     ) -> Self {
-        let prefix_nibbles = prefix.map(|p| Nibbles::from_bytes(p.as_bytes()).append_new(17));
+        let prefix_nibbles = prefix.map(|p| crate::trie_key::TrieKey::storage_prefix(p));
         Self {
             state_root,
             inner,
@@ -294,16 +294,14 @@ impl TrieWrapper {
     }
 }
 
-/// Prepends an account address prefix (with an invalid nibble `17` as separator) to a
-/// trie path, distinguishing storage trie entries from state trie entries in the flat
-/// key-value namespace. Returns the path unchanged if `prefix` is `None` (state trie).
+/// Prepends an account address prefix to a trie path, distinguishing storage trie
+/// entries from state trie entries in the flat key-value namespace.
+/// Returns the path unchanged if `prefix` is `None` (state trie).
+///
+/// Delegates to [`TrieKey::with_prefix`] — this function is kept for backward
+/// compatibility with call sites that work with raw `Nibbles`.
 pub fn apply_prefix(prefix: Option<H256>, path: Nibbles) -> Nibbles {
-    match prefix {
-        Some(prefix) => Nibbles::from_bytes(prefix.as_bytes())
-            .append_new(17)
-            .concat(&path),
-        None => path,
-    }
+    crate::trie_key::TrieKey::with_prefix(prefix, path).into_nibbles()
 }
 
 impl TrieDB for TrieWrapper {

crates/storage/lib.rs

@@ -71,13 +71,15 @@ mod layering;
 pub mod rlp;
 pub mod store;
 pub mod trie;
+pub mod trie_key;
 pub mod utils;
 
 pub use layering::apply_prefix;
 pub use store::{
     AccountUpdatesList, EngineType, Store, UpdateBatch, has_valid_db, hash_address, hash_key,
     read_chain_id_from_db,
 };
+pub use trie_key::TrieKey;
 
 /// Store Schema Version, must be updated on any breaking change.
 ///

crates/storage/store.rs

@@ -3072,8 +3072,11 @@ fn apply_trie_updates(
     let nodes = trie_mut.commit(root).unwrap_or_default();
     let mut result = Ok(());
     for (key, value) in nodes {
-        let is_leaf = key.len() == 65 || key.len() == 131;
-        let is_account = key.len() <= 65;
+        let trie_key = crate::trie_key::TrieKey::from_nibbles(
+            ethrex_trie::Nibbles::from_hex(key.clone()),
+        );
+        let is_leaf = trie_key.is_leaf();
+        let is_account = trie_key.is_account();
 
         if is_leaf && key > last_written {
             continue;
@@ -3244,7 +3247,7 @@ fn flatkeyvalue_generator(
                 write_txn.commit()?;
                 *last_computed_fkv
                     .write()
-                    .map_err(|_| StoreError::LockError)? = vec![0xff; 131];
+                    .map_err(|_| StoreError::LockError)? = vec![0xff; crate::trie_key::STORAGE_LEAF_LEN];
                 info!("FlatKeyValue generation finished.");
                 return Ok(());
             }

crates/storage/trie.rs

@@ -3,26 +3,25 @@ use crate::api::tables::{
 };
 use crate::api::{StorageBackend, StorageLockedView, StorageReadView};
 use crate::error::StoreError;
-use crate::layering::apply_prefix;
+use crate::trie_key::TrieKey;
 use ethrex_common::H256;
 use ethrex_trie::{Nibbles, TrieDB, error::TrieError};
 use std::sync::Arc;
 
 /// TrieDB implementation that holds a pre-acquired read view for the entire
 /// trie traversal, avoiding per-node-lookup allocation and lock acquisition.
+///
+/// Table routing (which RocksDB column family to read/write) is handled by
+/// [`TrieKey::table()`], based on the nibble-path length after prefix application.
 pub struct BackendTrieDB {
     /// Reference to the storage backend (used only for writes)
     db: Arc<dyn StorageBackend>,
     /// Pre-acquired read view held for the lifetime of this struct.
-    /// Using Arc allows sharing a single read view across multiple BackendTrieDB
-    /// instances (e.g., state trie + storage trie in a single query).
     read_view: Arc<dyn StorageReadView>,
     /// Last flatkeyvalue path already generated
     last_computed_flatkeyvalue: Nibbles,
-    nodes_table: &'static str,
-    fkv_table: &'static str,
-    /// Storage trie address prefix (for storage tries)
-    /// None for state tries, Some(address) for storage tries
+    /// Storage trie address prefix (for storage tries).
+    /// None for state tries, Some(keccak(address)) for storage tries.
     address_prefix: Option<H256>,
 }
 
@@ -47,8 +46,6 @@ impl BackendTrieDB {
             db,
             read_view,
             last_computed_flatkeyvalue,
-            nodes_table: ACCOUNT_TRIE_NODES,
-            fkv_table: ACCOUNT_FLATKEYVALUE,
             address_prefix: None,
         })
     }
@@ -73,8 +70,6 @@ impl BackendTrieDB {
             db,
             read_view,
             last_computed_flatkeyvalue,
-            nodes_table: STORAGE_TRIE_NODES,
-            fkv_table: STORAGE_FLATKEYVALUE,
             address_prefix: None,
         })
     }
@@ -101,38 +96,27 @@ impl BackendTrieDB {
             db,
             read_view,
             last_computed_flatkeyvalue,
-            nodes_table: STORAGE_TRIE_NODES,
-            fkv_table: STORAGE_FLATKEYVALUE,
             address_prefix: Some(address_prefix),
         })
     }
 
-    fn make_key(&self, path: Nibbles) -> Vec<u8> {
-        apply_prefix(self.address_prefix, path).into_vec()
-    }
-
-    /// Key might be for an account or storage slot
-    fn table_for_key(&self, key: &[u8]) -> &'static str {
-        let is_leaf = key.len() == 65 || key.len() == 131;
-        if is_leaf {
-            self.fkv_table
-        } else {
-            self.nodes_table
-        }
+    /// Build a [`TrieKey`] by applying this DB's address prefix to a raw path.
+    fn make_trie_key(&self, path: Nibbles) -> TrieKey {
+        TrieKey::with_prefix(self.address_prefix, path)
     }
 }
 
 impl TrieDB for BackendTrieDB {
     fn flatkeyvalue_computed(&self, key: Nibbles) -> bool {
-        let key = apply_prefix(self.address_prefix, key);
-        self.last_computed_flatkeyvalue >= key
+        let trie_key = self.make_trie_key(key);
+        self.last_computed_flatkeyvalue >= *trie_key.nibbles()
     }
 
     fn get(&self, key: Nibbles) -> Result<Option<Vec<u8>>, TrieError> {
-        let prefixed_key = self.make_key(key);
-        let table = self.table_for_key(&prefixed_key);
+        let trie_key = self.make_trie_key(key);
+        let table = trie_key.table();
         self.read_view
-            .get(table, prefixed_key.as_ref())
+            .get(table, trie_key.as_bytes())
             .map_err(|e| TrieError::DbError(anyhow::anyhow!("Failed to get from database: {}", e)))
     }
 
@@ -141,17 +125,20 @@ impl TrieDB for BackendTrieDB {
             TrieError::DbError(anyhow::anyhow!("Failed to begin write transaction: {}", e))
         })?;
         for (key, value) in key_values {
-            let prefixed_key = self.make_key(key);
-            let table = self.table_for_key(&prefixed_key);
-            tx.put_batch(table, vec![(prefixed_key, value)])
+            let trie_key = self.make_trie_key(key);
+            let table = trie_key.table();
+            tx.put_batch(table, vec![(trie_key.into_vec(), value)])
                 .map_err(|e| TrieError::DbError(anyhow::anyhow!("Failed to write batch: {}", e)))?;
         }
         tx.commit()
             .map_err(|e| TrieError::DbError(anyhow::anyhow!("Failed to write batch: {}", e)))
     }
 }
 
-/// Read-only version with persistent locked transaction/snapshot for batch reads
+/// Read-only version with persistent locked transaction/snapshot for batch reads.
+///
+/// Unlike [`BackendTrieDB`], this acquires per-CF locked snapshots at creation
+/// time, so that all reads within a batch see a consistent view.
 pub struct BackendTrieDBLocked {
     account_trie_tx: Box<dyn StorageLockedView>,
     storage_trie_tx: Box<dyn StorageLockedView>,
@@ -177,17 +164,16 @@ impl BackendTrieDBLocked {
         })
     }
 
-    /// Key is already prefixed
+    /// Select the locked view for the column family that this key belongs to.
     fn tx_for_key(&self, key: &Nibbles) -> &dyn StorageLockedView {
-        let is_leaf = key.len() == 65 || key.len() == 131;
-        let is_account = key.len() <= 65;
-        if is_leaf {
-            if is_account {
+        let trie_key = TrieKey::from_nibbles(key.clone());
+        if trie_key.is_leaf() {
+            if trie_key.is_account() {
                 &*self.account_fkv_tx
             } else {
                 &*self.storage_fkv_tx
             }
-        } else if is_account {
+        } else if trie_key.is_account() {
             &*self.account_trie_tx
         } else {
             &*self.storage_trie_tx

crates/storage/trie_key.rs

@@ -0,0 +1,154 @@
+use crate::api::tables::{
+    ACCOUNT_FLATKEYVALUE, ACCOUNT_TRIE_NODES, STORAGE_FLATKEYVALUE, STORAGE_TRIE_NODES,
+};
+use ethrex_common::{Address, H256};
+use ethrex_crypto::keccak::keccak_hash;
+use ethrex_trie::Nibbles;
+
+/// Nibble value that separates the account prefix from the storage path.
+///
+/// Value 17 is deliberately outside the valid nibble range (0–15), making it
+/// an unambiguous separator in the concatenated key.
+pub const SEPARATOR_NIBBLE: u8 = 17;
+
+/// Length (in nibbles) of an account leaf key:
+/// 64 nibbles from keccak(address) + 1 leaf-flag nibble.
+pub const ACCOUNT_LEAF_LEN: usize = 65;
+
+/// Length (in nibbles) of a storage leaf key:
+/// 65 (account prefix with separator) + 64 (keccak(slot)) + 2 (leaf flags) = 131.
+pub const STORAGE_LEAF_LEN: usize = 131;
+
+/// A fully-qualified key for the trie/FKV database.
+///
+/// Wraps a nibble path and provides methods for table routing and
+/// conversions. Constructed via typed constructors that handle hashing,
+/// nibble expansion, and prefix concatenation internally.
+#[derive(Debug, Clone)]
+pub struct TrieKey {
+    nibbles: Nibbles,
+}
+
+impl TrieKey {
+    // ── Constructors ─────────────────────────────────────────────
+
+    /// Hash an address and return its nibble path (64 nibbles + leaf flag = 65).
+    /// Used to look up an account in the state trie or FKV.
+    pub fn from_account_address(address: &Address) -> Self {
+        let hash = keccak_hash(address.to_fixed_bytes());
+        Self {
+            nibbles: Nibbles::from_bytes(&hash),
+        }
+    }
+
+    /// Build from a pre-computed account hash (H256).
+    pub fn from_account_hash(hash: H256) -> Self {
+        Self {
+            nibbles: Nibbles::from_bytes(hash.as_bytes()),
+        }
+    }
+
+    /// Build the prefix used for all storage keys of a given account:
+    /// `nibbles(keccak(address)) + [17]` (65 nibbles).
+    pub fn storage_prefix(address_hash: H256) -> Nibbles {
+        Nibbles::from_bytes(address_hash.as_bytes()).append_new(SEPARATOR_NIBBLE)
+    }
+
+    /// Build a full storage leaf key:
+    /// `nibbles(keccak(address)) + [17] + nibbles(keccak(slot))`.
+    pub fn from_storage_slot(address_hash: H256, slot: &H256) -> Self {
+        let prefix = Self::storage_prefix(address_hash);
+        let slot_hash = keccak_hash(slot.to_fixed_bytes());
+        let slot_nibbles = Nibbles::from_bytes(&slot_hash);
+        Self {
+            nibbles: prefix.concat(&slot_nibbles),
+        }
+    }
+
+    /// Wrap an existing nibble path as a TrieKey (for internal/intermediate nodes).
+    pub fn from_nibbles(nibbles: Nibbles) -> Self {
+        Self { nibbles }
+    }
+
+    /// Apply an optional account prefix to a raw trie path.
+    /// For account tries (prefix=None), returns the path unchanged.
+    /// For storage tries (prefix=Some(hash)), prepends the storage prefix.
+    pub fn with_prefix(prefix: Option<H256>, path: Nibbles) -> Self {
+        match prefix {
+            Some(hash) => {
+                let prefixed = Self::storage_prefix(hash).concat(&path);
+                Self { nibbles: prefixed }
+            }
+            None => Self { nibbles: path },
+        }
+    }
+
+    // ── Table routing ────────────────────────────────────────────
+
+    /// Determine which RocksDB column family this key belongs to.
+    pub fn table(&self) -> &'static str {
+        let len = self.nibbles.len();
+        let is_leaf = len == ACCOUNT_LEAF_LEN || len == STORAGE_LEAF_LEN;
+        let is_account = len <= ACCOUNT_LEAF_LEN;
+
+        if is_leaf {
+            if is_account {
+                ACCOUNT_FLATKEYVALUE
+            } else {
+                STORAGE_FLATKEYVALUE
+            }
+        } else if is_account {
+            ACCOUNT_TRIE_NODES
+        } else {
+            STORAGE_TRIE_NODES
+        }
+    }
+
+    /// Whether this key points to a leaf (account or storage value) vs an internal node.
+    pub fn is_leaf(&self) -> bool {
+        let len = self.nibbles.len();
+        len == ACCOUNT_LEAF_LEN || len == STORAGE_LEAF_LEN
+    }
+
+    /// Whether this key is in the account trie (vs storage trie).
+    pub fn is_account(&self) -> bool {
+        self.nibbles.len() <= ACCOUNT_LEAF_LEN
+    }
+
+    // ── Conversions ──────────────────────────────────────────────
+
+    /// Get the underlying Nibbles.
+    pub fn nibbles(&self) -> &Nibbles {
+        &self.nibbles
+    }
+
+    /// Consume and return the Nibbles.
+    pub fn into_nibbles(self) -> Nibbles {
+        self.nibbles
+    }
+
+    /// Convert to a byte vector (each nibble as one u8). Used as the RocksDB key.
+    pub fn into_vec(self) -> Vec<u8> {
+        self.nibbles.into_vec()
+    }
+
+    /// Borrow as byte slice (each nibble as one u8). Used for RocksDB lookups.
+    pub fn as_bytes(&self) -> &[u8] {
+        self.nibbles.as_ref()
+    }
+}
+
+// ── Standalone hash helpers ──────────────────────────────────────
+//
+// These are thin wrappers kept for call sites that only need the hash
+// (not a full TrieKey).
+
+/// Keccak-256 hash of an Ethereum address. Returns H256.
+pub fn hash_address(address: &Address) -> H256 {
+    H256(keccak_hash(address.to_fixed_bytes()))
+}
+
+/// Keccak-256 hash of a storage key. Returns the raw 32 bytes.
+pub fn hash_storage_key(key: &H256) -> [u8; 32] {
+    keccak_hash(key.to_fixed_bytes())
+}