1use crate::{DatabaseHashedCursorFactory, DatabaseTrieCursorFactory, TrieTableAdapter};
2use alloy_primitives::{keccak256, map::hash_map, Address, BlockNumber, B256};
3use reth_db_api::{models::BlockNumberAddress, transaction::DbTx};
4use reth_execution_errors::StorageRootError;
5use reth_storage_api::{BlockNumReader, StorageChangeSetReader};
6use reth_storage_errors::provider::ProviderResult;
7use reth_trie::{
8 hashed_cursor::HashedPostStateCursorFactory, trie_cursor::InMemoryTrieCursorFactory,
9 HashedStorage, StorageRoot, TrieInputSorted,
10};
11
12#[cfg(feature = "metrics")]
13use reth_trie::metrics::TrieRootMetrics;
14
15pub trait DatabaseStorageRoot<'a, TX> {
17 fn from_tx(tx: &'a TX, address: Address) -> Self;
19
20 fn from_tx_hashed(tx: &'a TX, hashed_address: B256) -> Self;
22
23 fn overlay_root(
25 tx: &'a TX,
26 address: Address,
27 input: TrieInputSorted,
28 ) -> Result<B256, StorageRootError>;
29}
30
31pub fn hashed_storage_from_reverts_with_provider<P>(
33 provider: &P,
34 address: Address,
35 from: BlockNumber,
36) -> ProviderResult<HashedStorage>
37where
38 P: StorageChangeSetReader + BlockNumReader,
39{
40 let mut storage = HashedStorage::default();
41 let tip = provider.last_block_number()?;
42
43 if from > tip {
44 return Ok(storage)
45 }
46
47 for (BlockNumberAddress((_, storage_address)), storage_change) in
48 provider.storage_changesets_range(from..=tip)?
49 {
50 if storage_address == address {
51 let hashed_slot = keccak256(storage_change.key);
52 if let hash_map::Entry::Vacant(entry) = storage.storage.entry(hashed_slot) {
53 entry.insert(storage_change.value);
54 }
55 }
56 }
57
58 Ok(storage)
59}
60
61impl<'a, TX: DbTx, A: TrieTableAdapter> DatabaseStorageRoot<'a, TX>
62 for StorageRoot<DatabaseTrieCursorFactory<&'a TX, A>, DatabaseHashedCursorFactory<&'a TX>>
63{
64 fn from_tx(tx: &'a TX, address: Address) -> Self {
65 Self::new(
66 DatabaseTrieCursorFactory::<_, A>::new(tx),
67 DatabaseHashedCursorFactory::new(tx),
68 address,
69 Default::default(),
70 #[cfg(feature = "metrics")]
71 TrieRootMetrics::new(reth_trie::TrieType::Storage),
72 )
73 }
74
75 fn from_tx_hashed(tx: &'a TX, hashed_address: B256) -> Self {
76 Self::new_hashed(
77 DatabaseTrieCursorFactory::<_, A>::new(tx),
78 DatabaseHashedCursorFactory::new(tx),
79 hashed_address,
80 Default::default(),
81 #[cfg(feature = "metrics")]
82 TrieRootMetrics::new(reth_trie::TrieType::Storage),
83 )
84 }
85
86 fn overlay_root(
87 tx: &'a TX,
88 address: Address,
89 mut input: TrieInputSorted,
90 ) -> Result<B256, StorageRootError> {
91 let hashed_address = keccak256(address);
92 StorageRoot::new(
93 InMemoryTrieCursorFactory::new(
94 DatabaseTrieCursorFactory::<_, A>::new(tx),
95 &input.nodes,
96 ),
97 HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &input.state),
98 address,
99 input
100 .prefix_sets
101 .storage_prefix_sets
102 .remove(&hashed_address)
103 .unwrap_or_default()
104 .freeze(),
105 #[cfg(feature = "metrics")]
106 TrieRootMetrics::new(reth_trie::TrieType::Storage),
107 )
108 .root()
109 }
110}
111
112#[cfg(test)]
113mod tests {
114 use super::*;
115 use alloy_consensus::Header;
116 use alloy_primitives::U256;
117 use reth_db_api::{models::BlockNumberAddress, tables, transaction::DbTxMut};
118 use reth_primitives_traits::StorageEntry;
119 use reth_provider::{
120 test_utils::create_test_provider_factory, StaticFileProviderFactory, StaticFileSegment,
121 StaticFileWriter, StorageSettingsCache,
122 };
123
124 fn append_storage_changesets_to_static_files(
125 factory: &impl StaticFileProviderFactory<
126 Primitives: reth_primitives_traits::NodePrimitives<BlockHeader = Header>,
127 >,
128 changesets: Vec<(u64, Vec<reth_db_api::models::StorageBeforeTx>)>,
129 ) {
130 let sf = factory.static_file_provider();
131 let mut writer = sf.latest_writer(StaticFileSegment::StorageChangeSets).unwrap();
132 for (block_number, changeset) in changesets {
133 writer.append_storage_changeset(changeset, block_number).unwrap();
134 }
135 writer.commit().unwrap();
136 }
137
138 fn append_headers_to_static_files(
139 factory: &impl StaticFileProviderFactory<
140 Primitives: reth_primitives_traits::NodePrimitives<BlockHeader = Header>,
141 >,
142 up_to_block: u64,
143 ) {
144 let sf = factory.static_file_provider();
145 let mut writer = sf.latest_writer(StaticFileSegment::Headers).unwrap();
146 let mut header = Header::default();
147 for num in 0..=up_to_block {
148 header.number = num;
149 writer.append_header(&header, &B256::ZERO).unwrap();
150 }
151 writer.commit().unwrap();
152 }
153
154 #[test]
155 fn test_hashed_storage_from_reverts_legacy() {
156 let factory = create_test_provider_factory();
157 let provider = factory.provider_rw().unwrap();
158
159 assert!(!provider.cached_storage_settings().use_hashed_state());
160
161 let address = Address::with_last_byte(42);
162 let slot1 = B256::from(U256::from(100));
163 let slot2 = B256::from(U256::from(200));
164
165 append_headers_to_static_files(&factory, 5);
166
167 provider
168 .tx_ref()
169 .put::<tables::StorageChangeSets>(
170 BlockNumberAddress((1, address)),
171 StorageEntry { key: slot1, value: U256::from(10) },
172 )
173 .unwrap();
174 provider
175 .tx_ref()
176 .put::<tables::StorageChangeSets>(
177 BlockNumberAddress((2, address)),
178 StorageEntry { key: slot2, value: U256::from(20) },
179 )
180 .unwrap();
181 provider
182 .tx_ref()
183 .put::<tables::StorageChangeSets>(
184 BlockNumberAddress((3, address)),
185 StorageEntry { key: slot1, value: U256::from(999) },
186 )
187 .unwrap();
188
189 let result = hashed_storage_from_reverts_with_provider(&*provider, address, 1).unwrap();
190
191 let hashed_slot1 = keccak256(slot1);
192 let hashed_slot2 = keccak256(slot2);
193
194 assert_eq!(result.storage.len(), 2);
195 assert_eq!(result.storage.get(&hashed_slot1), Some(&U256::from(10)));
196 assert_eq!(result.storage.get(&hashed_slot2), Some(&U256::from(20)));
197 }
198
199 #[test]
200 fn test_hashed_storage_from_reverts_hashed_state() {
201 use reth_db_api::models::{StorageBeforeTx, StorageSettings};
202
203 let factory = create_test_provider_factory();
204
205 factory.set_storage_settings_cache(StorageSettings::v2());
206
207 let provider = factory.provider_rw().unwrap();
208 assert!(provider.cached_storage_settings().use_hashed_state());
209 assert!(provider.cached_storage_settings().is_v2());
210
211 let address = Address::with_last_byte(42);
212 let plain_slot1 = B256::from(U256::from(100));
213 let plain_slot2 = B256::from(U256::from(200));
214 let hashed_slot1 = keccak256(plain_slot1);
215 let hashed_slot2 = keccak256(plain_slot2);
216
217 append_headers_to_static_files(&factory, 5);
218
219 append_storage_changesets_to_static_files(
220 &factory,
221 vec![
222 (0, vec![]),
223 (1, vec![StorageBeforeTx { address, key: plain_slot1, value: U256::from(10) }]),
224 (2, vec![StorageBeforeTx { address, key: plain_slot2, value: U256::from(20) }]),
225 (3, vec![StorageBeforeTx { address, key: plain_slot1, value: U256::from(999) }]),
226 ],
227 );
228
229 let result = hashed_storage_from_reverts_with_provider(&*provider, address, 1).unwrap();
230
231 assert_eq!(result.storage.len(), 2);
232 assert_eq!(result.storage.get(&hashed_slot1), Some(&U256::from(10)));
233 assert_eq!(result.storage.get(&hashed_slot2), Some(&U256::from(20)));
234 }
235}