From 56ed6a39387c6b485bc1351b8c5531ea47d4e801 Mon Sep 17 00:00:00 2001
From: James Prestwich <james@prestwi.ch>
Date: Fri, 6 Dec 2024 19:42:38 +0900
Subject: [PATCH] feat: couple convenience functions for nested cache dbs
 (#1852)

---
 crates/database/src/in_memory_db.rs | 36 +++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/crates/database/src/in_memory_db.rs b/crates/database/src/in_memory_db.rs
index a221c4f293..5d66a16597 100644
--- a/crates/database/src/in_memory_db.rs
+++ b/crates/database/src/in_memory_db.rs
@@ -38,7 +38,38 @@ impl<ExtDB: Default> Default for CacheDB<ExtDB> {
     }
 }
 
+impl<ExtDb> CacheDB<CacheDB<ExtDb>> {
+    /// Flatten a nested cache by applying the outer cache to the inner cache.
+    ///
+    /// The behavior is as follows:
+    /// - Accounts are overridden with outer accounts
+    /// - Contracts are overridden with outer contracts
+    /// - Logs are appended
+    /// - Block hashes are overridden with outer block hashes
+    pub fn flatten(self) -> CacheDB<ExtDb> {
+        let CacheDB {
+            accounts,
+            contracts,
+            logs,
+            block_hashes,
+            db: mut inner,
+        } = self;
+
+        inner.accounts.extend(accounts);
+        inner.contracts.extend(contracts);
+        inner.logs.extend(logs);
+        inner.block_hashes.extend(block_hashes);
+        inner
+    }
+
+    /// Discard the outer cache and return the inner cache.
+    pub fn discard_outer(self) -> CacheDB<ExtDb> {
+        self.db
+    }
+}
+
 impl<ExtDB> CacheDB<ExtDB> {
+    /// Create a new cache with the given external database.
     pub fn new(db: ExtDB) -> Self {
         let mut contracts = HashMap::default();
         contracts.insert(KECCAK_EMPTY, Bytecode::default());
@@ -78,6 +109,11 @@ impl<ExtDB> CacheDB<ExtDB> {
         self.insert_contract(&mut info);
         self.accounts.entry(address).or_default().info = info;
     }
+
+    /// Wrap the cache in a [CacheDB], creating a nested cache.
+    pub fn nest(self) -> CacheDB<Self> {
+        CacheDB::new(self)
+    }
 }
 
 impl<ExtDB: DatabaseRef> CacheDB<ExtDB> {