1use super::{
2 bundle_state::BundleRetention, cache::CacheState, plain_account::PlainStorage, BundleState,
3 CacheAccount, StateBuilder, TransitionAccount, TransitionState,
4};
5use bytecode::Bytecode;
6use database_interface::{Database, DatabaseCommit, DatabaseRef, EmptyDB};
7use primitives::{hash_map, Address, HashMap, StorageKey, StorageValue, B256, BLOCK_HASH_HISTORY};
8use state::{Account, AccountInfo};
9use std::{
10 boxed::Box,
11 collections::{btree_map, BTreeMap},
12 vec::Vec,
13};
14
15pub type DBBox<'a, E> = Box<dyn Database<Error = E> + Send + 'a>;
17
18pub type StateDBBox<'a, E> = State<DBBox<'a, E>>;
22
23#[derive(Debug)]
29pub struct State<DB> {
30 pub cache: CacheState,
37 pub database: DB,
43 pub transition_state: Option<TransitionState>,
47 pub bundle_state: BundleState,
53 pub use_preloaded_bundle: bool,
59 pub block_hashes: BTreeMap<u64, B256>,
66}
67
68impl State<EmptyDB> {
70 pub fn builder() -> StateBuilder<EmptyDB> {
72 StateBuilder::default()
73 }
74}
75
76impl<DB: Database> State<DB> {
77 pub fn bundle_size_hint(&self) -> usize {
81 self.bundle_state.size_hint()
82 }
83
84 pub fn increment_balances(
92 &mut self,
93 balances: impl IntoIterator<Item = (Address, u128)>,
94 ) -> Result<(), DB::Error> {
95 let mut transitions = Vec::new();
97 for (address, balance) in balances {
98 if balance == 0 {
99 continue;
100 }
101 let original_account = self.load_cache_account(address)?;
102 transitions.push((
103 address,
104 original_account
105 .increment_balance(balance)
106 .expect("Balance is not zero"),
107 ))
108 }
109 if let Some(s) = self.transition_state.as_mut() {
111 s.add_transitions(transitions)
112 }
113 Ok(())
114 }
115
116 pub fn drain_balances(
120 &mut self,
121 addresses: impl IntoIterator<Item = Address>,
122 ) -> Result<Vec<u128>, DB::Error> {
123 let mut transitions = Vec::new();
125 let mut balances = Vec::new();
126 for address in addresses {
127 let original_account = self.load_cache_account(address)?;
128 let (balance, transition) = original_account.drain_balance();
129 balances.push(balance);
130 transitions.push((address, transition))
131 }
132 if let Some(s) = self.transition_state.as_mut() {
134 s.add_transitions(transitions)
135 }
136 Ok(balances)
137 }
138
139 pub fn set_state_clear_flag(&mut self, has_state_clear: bool) {
141 self.cache.set_state_clear_flag(has_state_clear);
142 }
143
144 pub fn insert_not_existing(&mut self, address: Address) {
145 self.cache.insert_not_existing(address)
146 }
147
148 pub fn insert_account(&mut self, address: Address, info: AccountInfo) {
149 self.cache.insert_account(address, info)
150 }
151
152 pub fn insert_account_with_storage(
153 &mut self,
154 address: Address,
155 info: AccountInfo,
156 storage: PlainStorage,
157 ) {
158 self.cache
159 .insert_account_with_storage(address, info, storage)
160 }
161
162 pub fn apply_transition(&mut self, transitions: Vec<(Address, TransitionAccount)>) {
164 if let Some(s) = self.transition_state.as_mut() {
166 s.add_transitions(transitions)
167 }
168 }
169
170 pub fn merge_transitions(&mut self, retention: BundleRetention) {
176 if let Some(transition_state) = self.transition_state.as_mut().map(TransitionState::take) {
177 self.bundle_state
178 .apply_transitions_and_create_reverts(transition_state, retention);
179 }
180 }
181
182 pub fn load_cache_account(&mut self, address: Address) -> Result<&mut CacheAccount, DB::Error> {
187 match self.cache.accounts.entry(address) {
188 hash_map::Entry::Vacant(entry) => {
189 if self.use_preloaded_bundle {
190 if let Some(account) = self.bundle_state.account(&address).map(Into::into) {
192 return Ok(entry.insert(account));
193 }
194 }
195 let info = self.database.basic(address)?;
197 let account = match info {
198 None => CacheAccount::new_loaded_not_existing(),
199 Some(acc) if acc.is_empty() => {
200 CacheAccount::new_loaded_empty_eip161(HashMap::default())
201 }
202 Some(acc) => CacheAccount::new_loaded(acc, HashMap::default()),
203 };
204 Ok(entry.insert(account))
205 }
206 hash_map::Entry::Occupied(entry) => Ok(entry.into_mut()),
207 }
208 }
209
210 pub fn take_bundle(&mut self) -> BundleState {
222 core::mem::take(&mut self.bundle_state)
223 }
224}
225
226impl<DB: Database> Database for State<DB> {
227 type Error = DB::Error;
228
229 fn basic(&mut self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
230 self.load_cache_account(address).map(|a| a.account_info())
231 }
232
233 fn code_by_hash(&mut self, code_hash: B256) -> Result<Bytecode, Self::Error> {
234 let res = match self.cache.contracts.entry(code_hash) {
235 hash_map::Entry::Occupied(entry) => Ok(entry.get().clone()),
236 hash_map::Entry::Vacant(entry) => {
237 if self.use_preloaded_bundle {
238 if let Some(code) = self.bundle_state.contracts.get(&code_hash) {
239 entry.insert(code.clone());
240 return Ok(code.clone());
241 }
242 }
243 let code = self.database.code_by_hash(code_hash)?;
245 entry.insert(code.clone());
246 Ok(code)
247 }
248 };
249 res
250 }
251
252 fn storage(
253 &mut self,
254 address: Address,
255 index: StorageKey,
256 ) -> Result<StorageValue, Self::Error> {
257 if let Some(account) = self.cache.accounts.get_mut(&address) {
260 let is_storage_known = account.status.is_storage_known();
262 Ok(account
263 .account
264 .as_mut()
265 .map(|account| match account.storage.entry(index) {
266 hash_map::Entry::Occupied(entry) => Ok(*entry.get()),
267 hash_map::Entry::Vacant(entry) => {
268 let value = if is_storage_known {
271 StorageValue::ZERO
272 } else {
273 self.database.storage(address, index)?
274 };
275 entry.insert(value);
276 Ok(value)
277 }
278 })
279 .transpose()?
280 .unwrap_or_default())
281 } else {
282 unreachable!("For accessing any storage account is guaranteed to be loaded beforehand")
283 }
284 }
285
286 fn block_hash(&mut self, number: u64) -> Result<B256, Self::Error> {
287 match self.block_hashes.entry(number) {
288 btree_map::Entry::Occupied(entry) => Ok(*entry.get()),
289 btree_map::Entry::Vacant(entry) => {
290 let ret = *entry.insert(self.database.block_hash(number)?);
291
292 let last_block = number.saturating_sub(BLOCK_HASH_HISTORY);
294 while let Some(entry) = self.block_hashes.first_entry() {
295 if *entry.key() < last_block {
296 entry.remove();
297 } else {
298 break;
299 }
300 }
301
302 Ok(ret)
303 }
304 }
305 }
306}
307
308impl<DB: Database> DatabaseCommit for State<DB> {
309 fn commit(&mut self, evm_state: HashMap<Address, Account>) {
310 let transitions = self.cache.apply_evm_state(evm_state);
311 self.apply_transition(transitions);
312 }
313}
314
315impl<DB: DatabaseRef> DatabaseRef for State<DB> {
316 type Error = DB::Error;
317
318 fn basic_ref(&self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
319 if let Some(account) = self.cache.accounts.get(&address) {
321 return Ok(account.account_info());
322 }
323 if self.use_preloaded_bundle {
325 if let Some(account) = self.bundle_state.account(&address) {
326 return Ok(account.account_info());
327 }
328 }
329 self.database.basic_ref(address)
331 }
332
333 fn code_by_hash_ref(&self, code_hash: B256) -> Result<Bytecode, Self::Error> {
334 if let Some(code) = self.cache.contracts.get(&code_hash) {
336 return Ok(code.clone());
337 }
338 if self.use_preloaded_bundle {
340 if let Some(code) = self.bundle_state.contracts.get(&code_hash) {
341 return Ok(code.clone());
342 }
343 }
344 self.database.code_by_hash_ref(code_hash)
346 }
347
348 fn storage_ref(
349 &self,
350 address: Address,
351 index: StorageKey,
352 ) -> Result<StorageValue, Self::Error> {
353 if let Some(account) = self.cache.accounts.get(&address) {
355 if let Some(plain_account) = &account.account {
356 if let Some(storage_value) = plain_account.storage.get(&index) {
358 return Ok(*storage_value);
359 }
360 if account.status.is_storage_known() {
363 return Ok(StorageValue::ZERO);
364 }
365 }
366 }
367 self.database.storage_ref(address, index)
369 }
370
371 fn block_hash_ref(&self, number: u64) -> Result<B256, Self::Error> {
372 if let Some(entry) = self.block_hashes.get(&number) {
373 return Ok(*entry);
374 }
375 self.database.block_hash_ref(number)
377 }
378}
379
380#[cfg(test)]
381mod tests {
382 use super::*;
383 use crate::{
384 states::{reverts::AccountInfoRevert, StorageSlot},
385 AccountRevert, AccountStatus, BundleAccount, RevertToSlot,
386 };
387 use primitives::{keccak256, U256};
388
389 #[test]
390 fn block_hash_cache() {
391 let mut state = State::builder().build();
392 state.block_hash(1u64).unwrap();
393 state.block_hash(2u64).unwrap();
394
395 let test_number = BLOCK_HASH_HISTORY + 2;
396
397 let block1_hash = keccak256(U256::from(1).to_string().as_bytes());
398 let block2_hash = keccak256(U256::from(2).to_string().as_bytes());
399 let block_test_hash = keccak256(U256::from(test_number).to_string().as_bytes());
400
401 assert_eq!(
402 state.block_hashes,
403 BTreeMap::from([(1, block1_hash), (2, block2_hash)])
404 );
405
406 state.block_hash(test_number).unwrap();
407 assert_eq!(
408 state.block_hashes,
409 BTreeMap::from([(test_number, block_test_hash), (2, block2_hash)])
410 );
411 }
412
413 #[test]
420 fn reverts_preserve_old_values() {
421 let mut state = State::builder().with_bundle_update().build();
422
423 let (slot1, slot2, slot3) = (
424 StorageKey::from(1),
425 StorageKey::from(2),
426 StorageKey::from(3),
427 );
428
429 let new_account_address = Address::from_slice(&[0x1; 20]);
432 let new_account_created_info = AccountInfo {
433 nonce: 1,
434 balance: U256::from(1),
435 ..Default::default()
436 };
437 let new_account_changed_info = AccountInfo {
438 nonce: 2,
439 ..new_account_created_info.clone()
440 };
441 let new_account_changed_info2 = AccountInfo {
442 nonce: 3,
443 ..new_account_changed_info.clone()
444 };
445
446 let existing_account_address = Address::from_slice(&[0x2; 20]);
448 let existing_account_initial_info = AccountInfo {
449 nonce: 1,
450 ..Default::default()
451 };
452 let existing_account_initial_storage = HashMap::<StorageKey, StorageValue>::from_iter([
453 (slot1, StorageValue::from(100)), (slot2, StorageValue::from(200)), ]);
456 let existing_account_changed_info = AccountInfo {
457 nonce: 2,
458 ..existing_account_initial_info.clone()
459 };
460
461 state.apply_transition(Vec::from([
463 (
464 new_account_address,
465 TransitionAccount {
466 status: AccountStatus::InMemoryChange,
467 info: Some(new_account_created_info.clone()),
468 previous_status: AccountStatus::LoadedNotExisting,
469 previous_info: None,
470 ..Default::default()
471 },
472 ),
473 (
474 existing_account_address,
475 TransitionAccount {
476 status: AccountStatus::InMemoryChange,
477 info: Some(existing_account_changed_info.clone()),
478 previous_status: AccountStatus::Loaded,
479 previous_info: Some(existing_account_initial_info.clone()),
480 storage: HashMap::from_iter([(
481 slot1,
482 StorageSlot::new_changed(
483 *existing_account_initial_storage.get(&slot1).unwrap(),
484 StorageValue::from(1000),
485 ),
486 )]),
487 storage_was_destroyed: false,
488 },
489 ),
490 ]));
491
492 state.apply_transition(Vec::from([(
494 new_account_address,
495 TransitionAccount {
496 status: AccountStatus::InMemoryChange,
497 info: Some(new_account_changed_info.clone()),
498 previous_status: AccountStatus::InMemoryChange,
499 previous_info: Some(new_account_created_info.clone()),
500 ..Default::default()
501 },
502 )]));
503
504 state.apply_transition(Vec::from([
506 (
507 new_account_address,
508 TransitionAccount {
509 status: AccountStatus::InMemoryChange,
510 info: Some(new_account_changed_info2.clone()),
511 previous_status: AccountStatus::InMemoryChange,
512 previous_info: Some(new_account_changed_info),
513 storage: HashMap::from_iter([(
514 slot1,
515 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(1)),
516 )]),
517 storage_was_destroyed: false,
518 },
519 ),
520 (
521 existing_account_address,
522 TransitionAccount {
523 status: AccountStatus::InMemoryChange,
524 info: Some(existing_account_changed_info.clone()),
525 previous_status: AccountStatus::InMemoryChange,
526 previous_info: Some(existing_account_changed_info.clone()),
527 storage: HashMap::from_iter([
528 (
529 slot1,
530 StorageSlot::new_changed(
531 StorageValue::from(100),
532 StorageValue::from(1_000),
533 ),
534 ),
535 (
536 slot2,
537 StorageSlot::new_changed(
538 *existing_account_initial_storage.get(&slot2).unwrap(),
539 StorageValue::from(2_000),
540 ),
541 ),
542 (
544 slot3,
545 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(3_000)),
546 ),
547 ]),
548 storage_was_destroyed: false,
549 },
550 ),
551 ]));
552
553 state.merge_transitions(BundleRetention::Reverts);
554 let mut bundle_state = state.take_bundle();
555
556 bundle_state.reverts.sort();
559 assert_eq!(
560 bundle_state.reverts.as_ref(),
561 Vec::from([Vec::from([
562 (
563 new_account_address,
564 AccountRevert {
565 account: AccountInfoRevert::DeleteIt,
566 previous_status: AccountStatus::LoadedNotExisting,
567 storage: HashMap::from_iter([(
568 slot1,
569 RevertToSlot::Some(StorageValue::ZERO)
570 )]),
571 wipe_storage: false,
572 }
573 ),
574 (
575 existing_account_address,
576 AccountRevert {
577 account: AccountInfoRevert::RevertTo(existing_account_initial_info.clone()),
578 previous_status: AccountStatus::Loaded,
579 storage: HashMap::from_iter([
580 (
581 slot1,
582 RevertToSlot::Some(
583 *existing_account_initial_storage.get(&slot1).unwrap()
584 )
585 ),
586 (
587 slot2,
588 RevertToSlot::Some(
589 *existing_account_initial_storage.get(&slot2).unwrap()
590 )
591 ),
592 (slot3, RevertToSlot::Some(StorageValue::ZERO))
593 ]),
594 wipe_storage: false,
595 }
596 ),
597 ])]),
598 "The account or storage reverts are incorrect"
599 );
600
601 assert_eq!(
604 bundle_state.account(&new_account_address),
605 Some(&BundleAccount {
606 info: Some(new_account_changed_info2),
607 original_info: None,
608 status: AccountStatus::InMemoryChange,
609 storage: HashMap::from_iter([(
610 slot1,
611 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(1))
612 )]),
613 }),
614 "The latest state of the new account is incorrect"
615 );
616
617 assert_eq!(
620 bundle_state.account(&existing_account_address),
621 Some(&BundleAccount {
622 info: Some(existing_account_changed_info),
623 original_info: Some(existing_account_initial_info),
624 status: AccountStatus::InMemoryChange,
625 storage: HashMap::from_iter([
626 (
627 slot1,
628 StorageSlot::new_changed(
629 *existing_account_initial_storage.get(&slot1).unwrap(),
630 StorageValue::from(1_000)
631 )
632 ),
633 (
634 slot2,
635 StorageSlot::new_changed(
636 *existing_account_initial_storage.get(&slot2).unwrap(),
637 StorageValue::from(2_000)
638 )
639 ),
640 (
642 slot3,
643 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(3_000))
644 ),
645 ]),
646 }),
647 "The latest state of the existing account is incorrect"
648 );
649 }
650
651 #[test]
654 fn bundle_scoped_reverts_collapse() {
655 let mut state = State::builder().with_bundle_update().build();
656
657 let new_account_address = Address::from_slice(&[0x1; 20]);
659 let new_account_created_info = AccountInfo {
660 nonce: 1,
661 balance: U256::from(1),
662 ..Default::default()
663 };
664
665 let existing_account_address = Address::from_slice(&[0x2; 20]);
667 let existing_account_initial_info = AccountInfo {
668 nonce: 1,
669 ..Default::default()
670 };
671 let existing_account_updated_info = AccountInfo {
672 nonce: 1,
673 balance: U256::from(1),
674 ..Default::default()
675 };
676
677 let (slot1, slot2) = (StorageKey::from(1), StorageKey::from(2));
679 let existing_account_with_storage_address = Address::from_slice(&[0x3; 20]);
680 let existing_account_with_storage_info = AccountInfo {
681 nonce: 1,
682 ..Default::default()
683 };
684 state.apply_transition(Vec::from([
686 (
687 new_account_address,
688 TransitionAccount {
689 status: AccountStatus::InMemoryChange,
690 info: Some(new_account_created_info.clone()),
691 previous_status: AccountStatus::LoadedNotExisting,
692 previous_info: None,
693 ..Default::default()
694 },
695 ),
696 (
697 existing_account_address,
698 TransitionAccount {
699 status: AccountStatus::Changed,
700 info: Some(existing_account_updated_info.clone()),
701 previous_status: AccountStatus::Loaded,
702 previous_info: Some(existing_account_initial_info.clone()),
703 ..Default::default()
704 },
705 ),
706 (
707 existing_account_with_storage_address,
708 TransitionAccount {
709 status: AccountStatus::Changed,
710 info: Some(existing_account_with_storage_info.clone()),
711 previous_status: AccountStatus::Loaded,
712 previous_info: Some(existing_account_with_storage_info.clone()),
713 storage: HashMap::from_iter([
714 (
715 slot1,
716 StorageSlot::new_changed(StorageValue::from(1), StorageValue::from(10)),
717 ),
718 (
719 slot2,
720 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(20)),
721 ),
722 ]),
723 storage_was_destroyed: false,
724 },
725 ),
726 ]));
727
728 state.apply_transition(Vec::from([
730 (
731 new_account_address,
732 TransitionAccount {
733 status: AccountStatus::Destroyed,
734 info: None,
735 previous_status: AccountStatus::InMemoryChange,
736 previous_info: Some(new_account_created_info),
737 ..Default::default()
738 },
739 ),
740 (
741 existing_account_address,
742 TransitionAccount {
743 status: AccountStatus::Changed,
744 info: Some(existing_account_initial_info),
745 previous_status: AccountStatus::Changed,
746 previous_info: Some(existing_account_updated_info),
747 ..Default::default()
748 },
749 ),
750 (
751 existing_account_with_storage_address,
752 TransitionAccount {
753 status: AccountStatus::Changed,
754 info: Some(existing_account_with_storage_info.clone()),
755 previous_status: AccountStatus::Changed,
756 previous_info: Some(existing_account_with_storage_info.clone()),
757 storage: HashMap::from_iter([
758 (
759 slot1,
760 StorageSlot::new_changed(StorageValue::from(10), StorageValue::from(1)),
761 ),
762 (
763 slot2,
764 StorageSlot::new_changed(StorageValue::from(20), StorageValue::ZERO),
765 ),
766 ]),
767 storage_was_destroyed: false,
768 },
769 ),
770 ]));
771
772 state.merge_transitions(BundleRetention::Reverts);
773
774 let mut bundle_state = state.take_bundle();
775 bundle_state.reverts.sort();
776
777 assert_eq!(bundle_state.reverts.as_ref(), Vec::from([Vec::from([])]));
780 }
781
782 #[test]
784 fn selfdestruct_state_and_reverts() {
785 let mut state = State::builder().with_bundle_update().build();
786
787 let existing_account_address = Address::from_slice(&[0x1; 20]);
789 let existing_account_info = AccountInfo {
790 nonce: 1,
791 ..Default::default()
792 };
793
794 let (slot1, slot2) = (StorageKey::from(1), StorageKey::from(2));
795
796 state.apply_transition(Vec::from([(
798 existing_account_address,
799 TransitionAccount {
800 status: AccountStatus::Destroyed,
801 info: None,
802 previous_status: AccountStatus::Loaded,
803 previous_info: Some(existing_account_info.clone()),
804 storage: HashMap::default(),
805 storage_was_destroyed: true,
806 },
807 )]));
808
809 state.apply_transition(Vec::from([(
811 existing_account_address,
812 TransitionAccount {
813 status: AccountStatus::DestroyedChanged,
814 info: Some(existing_account_info.clone()),
815 previous_status: AccountStatus::Destroyed,
816 previous_info: None,
817 storage: HashMap::from_iter([(
818 slot1,
819 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(1)),
820 )]),
821 storage_was_destroyed: false,
822 },
823 )]));
824
825 state.apply_transition(Vec::from([(
827 existing_account_address,
828 TransitionAccount {
829 status: AccountStatus::DestroyedAgain,
830 info: None,
831 previous_status: AccountStatus::DestroyedChanged,
832 previous_info: Some(existing_account_info.clone()),
833 storage: HashMap::default(),
835 storage_was_destroyed: true,
836 },
837 )]));
838
839 state.apply_transition(Vec::from([(
841 existing_account_address,
842 TransitionAccount {
843 status: AccountStatus::DestroyedChanged,
844 info: Some(existing_account_info.clone()),
845 previous_status: AccountStatus::DestroyedAgain,
846 previous_info: None,
847 storage: HashMap::from_iter([(
848 slot2,
849 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(2)),
850 )]),
851 storage_was_destroyed: false,
852 },
853 )]));
854
855 state.merge_transitions(BundleRetention::Reverts);
856
857 let bundle_state = state.take_bundle();
858
859 assert_eq!(
860 bundle_state.state,
861 HashMap::from_iter([(
862 existing_account_address,
863 BundleAccount {
864 info: Some(existing_account_info.clone()),
865 original_info: Some(existing_account_info.clone()),
866 storage: HashMap::from_iter([(
867 slot2,
868 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(2))
869 )]),
870 status: AccountStatus::DestroyedChanged,
871 }
872 )])
873 );
874
875 assert_eq!(
876 bundle_state.reverts.as_ref(),
877 Vec::from([Vec::from([(
878 existing_account_address,
879 AccountRevert {
880 account: AccountInfoRevert::DoNothing,
881 previous_status: AccountStatus::Loaded,
882 storage: HashMap::from_iter([(slot2, RevertToSlot::Destroyed)]),
883 wipe_storage: true,
884 }
885 )])])
886 )
887 }
888}