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 balances = balances.into_iter();
97 let mut transitions = Vec::with_capacity(balances.size_hint().0);
98 for (address, balance) in balances {
99 if balance == 0 {
100 continue;
101 }
102 let original_account = self.load_cache_account(address)?;
103 transitions.push((
104 address,
105 original_account
106 .increment_balance(balance)
107 .expect("Balance is not zero"),
108 ))
109 }
110 if let Some(s) = self.transition_state.as_mut() {
112 s.add_transitions(transitions)
113 }
114 Ok(())
115 }
116
117 pub fn drain_balances(
121 &mut self,
122 addresses: impl IntoIterator<Item = Address>,
123 ) -> Result<Vec<u128>, DB::Error> {
124 let mut transitions = Vec::new();
126 let mut balances = Vec::new();
127 for address in addresses {
128 let original_account = self.load_cache_account(address)?;
129 let (balance, transition) = original_account.drain_balance();
130 balances.push(balance);
131 transitions.push((address, transition))
132 }
133 if let Some(s) = self.transition_state.as_mut() {
135 s.add_transitions(transitions)
136 }
137 Ok(balances)
138 }
139
140 pub fn set_state_clear_flag(&mut self, has_state_clear: bool) {
142 self.cache.set_state_clear_flag(has_state_clear);
143 }
144
145 pub fn insert_not_existing(&mut self, address: Address) {
147 self.cache.insert_not_existing(address)
148 }
149
150 pub fn insert_account(&mut self, address: Address, info: AccountInfo) {
152 self.cache.insert_account(address, info)
153 }
154
155 pub fn insert_account_with_storage(
157 &mut self,
158 address: Address,
159 info: AccountInfo,
160 storage: PlainStorage,
161 ) {
162 self.cache
163 .insert_account_with_storage(address, info, storage)
164 }
165
166 pub fn apply_transition(&mut self, transitions: Vec<(Address, TransitionAccount)>) {
168 if let Some(s) = self.transition_state.as_mut() {
170 s.add_transitions(transitions)
171 }
172 }
173
174 pub fn merge_transitions(&mut self, retention: BundleRetention) {
180 if let Some(transition_state) = self.transition_state.as_mut().map(TransitionState::take) {
181 self.bundle_state
182 .apply_transitions_and_create_reverts(transition_state, retention);
183 }
184 }
185
186 pub fn load_cache_account(&mut self, address: Address) -> Result<&mut CacheAccount, DB::Error> {
191 match self.cache.accounts.entry(address) {
192 hash_map::Entry::Vacant(entry) => {
193 if self.use_preloaded_bundle {
194 if let Some(account) = self.bundle_state.account(&address).map(Into::into) {
196 return Ok(entry.insert(account));
197 }
198 }
199 let info = self.database.basic(address)?;
201 let account = match info {
202 None => CacheAccount::new_loaded_not_existing(),
203 Some(acc) if acc.is_empty() => {
204 CacheAccount::new_loaded_empty_eip161(HashMap::default())
205 }
206 Some(acc) => CacheAccount::new_loaded(acc, HashMap::default()),
207 };
208 Ok(entry.insert(account))
209 }
210 hash_map::Entry::Occupied(entry) => Ok(entry.into_mut()),
211 }
212 }
213
214 pub fn take_bundle(&mut self) -> BundleState {
226 core::mem::take(&mut self.bundle_state)
227 }
228}
229
230impl<DB: Database> Database for State<DB> {
231 type Error = DB::Error;
232
233 fn basic(&mut self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
234 self.load_cache_account(address).map(|a| a.account_info())
235 }
236
237 fn code_by_hash(&mut self, code_hash: B256) -> Result<Bytecode, Self::Error> {
238 let res = match self.cache.contracts.entry(code_hash) {
239 hash_map::Entry::Occupied(entry) => Ok(entry.get().clone()),
240 hash_map::Entry::Vacant(entry) => {
241 if self.use_preloaded_bundle {
242 if let Some(code) = self.bundle_state.contracts.get(&code_hash) {
243 entry.insert(code.clone());
244 return Ok(code.clone());
245 }
246 }
247 let code = self.database.code_by_hash(code_hash)?;
249 entry.insert(code.clone());
250 Ok(code)
251 }
252 };
253 res
254 }
255
256 fn storage(
257 &mut self,
258 address: Address,
259 index: StorageKey,
260 ) -> Result<StorageValue, Self::Error> {
261 let account = if let Some(account) = self.cache.accounts.get_mut(&address) {
263 account
264 } else {
265 self.load_cache_account(address)?;
266 self.cache.accounts.get_mut(&address).unwrap()
268 };
269
270 let is_storage_known = account.status.is_storage_known();
272 Ok(account
273 .account
274 .as_mut()
275 .map(|account| match account.storage.entry(index) {
276 hash_map::Entry::Occupied(entry) => Ok(*entry.get()),
277 hash_map::Entry::Vacant(entry) => {
278 let value = if is_storage_known {
281 StorageValue::ZERO
282 } else {
283 self.database.storage(address, index)?
284 };
285 entry.insert(value);
286 Ok(value)
287 }
288 })
289 .transpose()?
290 .unwrap_or_default())
291 }
292
293 fn block_hash(&mut self, number: u64) -> Result<B256, Self::Error> {
294 match self.block_hashes.entry(number) {
295 btree_map::Entry::Occupied(entry) => Ok(*entry.get()),
296 btree_map::Entry::Vacant(entry) => {
297 let ret = *entry.insert(self.database.block_hash(number)?);
298
299 let last_block = number.saturating_sub(BLOCK_HASH_HISTORY);
301 while let Some(entry) = self.block_hashes.first_entry() {
302 if *entry.key() < last_block {
303 entry.remove();
304 } else {
305 break;
306 }
307 }
308
309 Ok(ret)
310 }
311 }
312 }
313}
314
315impl<DB: Database> DatabaseCommit for State<DB> {
316 fn commit(&mut self, evm_state: HashMap<Address, Account>) {
317 let transitions = self.cache.apply_evm_state(evm_state);
318 self.apply_transition(transitions);
319 }
320}
321
322impl<DB: DatabaseRef> DatabaseRef for State<DB> {
323 type Error = DB::Error;
324
325 fn basic_ref(&self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
326 if let Some(account) = self.cache.accounts.get(&address) {
328 return Ok(account.account_info());
329 }
330 if self.use_preloaded_bundle {
332 if let Some(account) = self.bundle_state.account(&address) {
333 return Ok(account.account_info());
334 }
335 }
336 self.database.basic_ref(address)
338 }
339
340 fn code_by_hash_ref(&self, code_hash: B256) -> Result<Bytecode, Self::Error> {
341 if let Some(code) = self.cache.contracts.get(&code_hash) {
343 return Ok(code.clone());
344 }
345 if self.use_preloaded_bundle {
347 if let Some(code) = self.bundle_state.contracts.get(&code_hash) {
348 return Ok(code.clone());
349 }
350 }
351 self.database.code_by_hash_ref(code_hash)
353 }
354
355 fn storage_ref(
356 &self,
357 address: Address,
358 index: StorageKey,
359 ) -> Result<StorageValue, Self::Error> {
360 if let Some(account) = self.cache.accounts.get(&address) {
362 if let Some(plain_account) = &account.account {
363 if let Some(storage_value) = plain_account.storage.get(&index) {
365 return Ok(*storage_value);
366 }
367 if account.status.is_storage_known() {
370 return Ok(StorageValue::ZERO);
371 }
372 }
373 }
374 self.database.storage_ref(address, index)
376 }
377
378 fn block_hash_ref(&self, number: u64) -> Result<B256, Self::Error> {
379 if let Some(entry) = self.block_hashes.get(&number) {
380 return Ok(*entry);
381 }
382 self.database.block_hash_ref(number)
384 }
385}
386
387#[cfg(test)]
388mod tests {
389 use super::*;
390 use crate::{
391 states::{reverts::AccountInfoRevert, StorageSlot},
392 AccountRevert, AccountStatus, BundleAccount, RevertToSlot,
393 };
394 use primitives::{keccak256, U256};
395
396 #[test]
397 fn block_hash_cache() {
398 let mut state = State::builder().build();
399 state.block_hash(1u64).unwrap();
400 state.block_hash(2u64).unwrap();
401
402 let test_number = BLOCK_HASH_HISTORY + 2;
403
404 let block1_hash = keccak256(U256::from(1).to_string().as_bytes());
405 let block2_hash = keccak256(U256::from(2).to_string().as_bytes());
406 let block_test_hash = keccak256(U256::from(test_number).to_string().as_bytes());
407
408 assert_eq!(
409 state.block_hashes,
410 BTreeMap::from([(1, block1_hash), (2, block2_hash)])
411 );
412
413 state.block_hash(test_number).unwrap();
414 assert_eq!(
415 state.block_hashes,
416 BTreeMap::from([(test_number, block_test_hash), (2, block2_hash)])
417 );
418 }
419
420 #[test]
427 fn reverts_preserve_old_values() {
428 let mut state = State::builder().with_bundle_update().build();
429
430 let (slot1, slot2, slot3) = (
431 StorageKey::from(1),
432 StorageKey::from(2),
433 StorageKey::from(3),
434 );
435
436 let new_account_address = Address::from_slice(&[0x1; 20]);
439 let new_account_created_info = AccountInfo {
440 nonce: 1,
441 balance: U256::from(1),
442 ..Default::default()
443 };
444 let new_account_changed_info = AccountInfo {
445 nonce: 2,
446 ..new_account_created_info.clone()
447 };
448 let new_account_changed_info2 = AccountInfo {
449 nonce: 3,
450 ..new_account_changed_info.clone()
451 };
452
453 let existing_account_address = Address::from_slice(&[0x2; 20]);
455 let existing_account_initial_info = AccountInfo {
456 nonce: 1,
457 ..Default::default()
458 };
459 let existing_account_initial_storage = HashMap::<StorageKey, StorageValue>::from_iter([
460 (slot1, StorageValue::from(100)), (slot2, StorageValue::from(200)), ]);
463 let existing_account_changed_info = AccountInfo {
464 nonce: 2,
465 ..existing_account_initial_info.clone()
466 };
467
468 state.apply_transition(Vec::from([
470 (
471 new_account_address,
472 TransitionAccount {
473 status: AccountStatus::InMemoryChange,
474 info: Some(new_account_created_info.clone()),
475 previous_status: AccountStatus::LoadedNotExisting,
476 previous_info: None,
477 ..Default::default()
478 },
479 ),
480 (
481 existing_account_address,
482 TransitionAccount {
483 status: AccountStatus::InMemoryChange,
484 info: Some(existing_account_changed_info.clone()),
485 previous_status: AccountStatus::Loaded,
486 previous_info: Some(existing_account_initial_info.clone()),
487 storage: HashMap::from_iter([(
488 slot1,
489 StorageSlot::new_changed(
490 *existing_account_initial_storage.get(&slot1).unwrap(),
491 StorageValue::from(1000),
492 ),
493 )]),
494 storage_was_destroyed: false,
495 },
496 ),
497 ]));
498
499 state.apply_transition(Vec::from([(
501 new_account_address,
502 TransitionAccount {
503 status: AccountStatus::InMemoryChange,
504 info: Some(new_account_changed_info.clone()),
505 previous_status: AccountStatus::InMemoryChange,
506 previous_info: Some(new_account_created_info.clone()),
507 ..Default::default()
508 },
509 )]));
510
511 state.apply_transition(Vec::from([
513 (
514 new_account_address,
515 TransitionAccount {
516 status: AccountStatus::InMemoryChange,
517 info: Some(new_account_changed_info2.clone()),
518 previous_status: AccountStatus::InMemoryChange,
519 previous_info: Some(new_account_changed_info),
520 storage: HashMap::from_iter([(
521 slot1,
522 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(1)),
523 )]),
524 storage_was_destroyed: false,
525 },
526 ),
527 (
528 existing_account_address,
529 TransitionAccount {
530 status: AccountStatus::InMemoryChange,
531 info: Some(existing_account_changed_info.clone()),
532 previous_status: AccountStatus::InMemoryChange,
533 previous_info: Some(existing_account_changed_info.clone()),
534 storage: HashMap::from_iter([
535 (
536 slot1,
537 StorageSlot::new_changed(
538 StorageValue::from(100),
539 StorageValue::from(1_000),
540 ),
541 ),
542 (
543 slot2,
544 StorageSlot::new_changed(
545 *existing_account_initial_storage.get(&slot2).unwrap(),
546 StorageValue::from(2_000),
547 ),
548 ),
549 (
551 slot3,
552 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(3_000)),
553 ),
554 ]),
555 storage_was_destroyed: false,
556 },
557 ),
558 ]));
559
560 state.merge_transitions(BundleRetention::Reverts);
561 let mut bundle_state = state.take_bundle();
562
563 bundle_state.reverts.sort();
566 assert_eq!(
567 bundle_state.reverts.as_ref(),
568 Vec::from([Vec::from([
569 (
570 new_account_address,
571 AccountRevert {
572 account: AccountInfoRevert::DeleteIt,
573 previous_status: AccountStatus::LoadedNotExisting,
574 storage: HashMap::from_iter([(
575 slot1,
576 RevertToSlot::Some(StorageValue::ZERO)
577 )]),
578 wipe_storage: false,
579 }
580 ),
581 (
582 existing_account_address,
583 AccountRevert {
584 account: AccountInfoRevert::RevertTo(existing_account_initial_info.clone()),
585 previous_status: AccountStatus::Loaded,
586 storage: HashMap::from_iter([
587 (
588 slot1,
589 RevertToSlot::Some(
590 *existing_account_initial_storage.get(&slot1).unwrap()
591 )
592 ),
593 (
594 slot2,
595 RevertToSlot::Some(
596 *existing_account_initial_storage.get(&slot2).unwrap()
597 )
598 ),
599 (slot3, RevertToSlot::Some(StorageValue::ZERO))
600 ]),
601 wipe_storage: false,
602 }
603 ),
604 ])]),
605 "The account or storage reverts are incorrect"
606 );
607
608 assert_eq!(
611 bundle_state.account(&new_account_address),
612 Some(&BundleAccount {
613 info: Some(new_account_changed_info2),
614 original_info: None,
615 status: AccountStatus::InMemoryChange,
616 storage: HashMap::from_iter([(
617 slot1,
618 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(1))
619 )]),
620 }),
621 "The latest state of the new account is incorrect"
622 );
623
624 assert_eq!(
627 bundle_state.account(&existing_account_address),
628 Some(&BundleAccount {
629 info: Some(existing_account_changed_info),
630 original_info: Some(existing_account_initial_info),
631 status: AccountStatus::InMemoryChange,
632 storage: HashMap::from_iter([
633 (
634 slot1,
635 StorageSlot::new_changed(
636 *existing_account_initial_storage.get(&slot1).unwrap(),
637 StorageValue::from(1_000)
638 )
639 ),
640 (
641 slot2,
642 StorageSlot::new_changed(
643 *existing_account_initial_storage.get(&slot2).unwrap(),
644 StorageValue::from(2_000)
645 )
646 ),
647 (
649 slot3,
650 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(3_000))
651 ),
652 ]),
653 }),
654 "The latest state of the existing account is incorrect"
655 );
656 }
657
658 #[test]
661 fn bundle_scoped_reverts_collapse() {
662 let mut state = State::builder().with_bundle_update().build();
663
664 let new_account_address = Address::from_slice(&[0x1; 20]);
666 let new_account_created_info = AccountInfo {
667 nonce: 1,
668 balance: U256::from(1),
669 ..Default::default()
670 };
671
672 let existing_account_address = Address::from_slice(&[0x2; 20]);
674 let existing_account_initial_info = AccountInfo {
675 nonce: 1,
676 ..Default::default()
677 };
678 let existing_account_updated_info = AccountInfo {
679 nonce: 1,
680 balance: U256::from(1),
681 ..Default::default()
682 };
683
684 let (slot1, slot2) = (StorageKey::from(1), StorageKey::from(2));
686 let existing_account_with_storage_address = Address::from_slice(&[0x3; 20]);
687 let existing_account_with_storage_info = AccountInfo {
688 nonce: 1,
689 ..Default::default()
690 };
691 state.apply_transition(Vec::from([
693 (
694 new_account_address,
695 TransitionAccount {
696 status: AccountStatus::InMemoryChange,
697 info: Some(new_account_created_info.clone()),
698 previous_status: AccountStatus::LoadedNotExisting,
699 previous_info: None,
700 ..Default::default()
701 },
702 ),
703 (
704 existing_account_address,
705 TransitionAccount {
706 status: AccountStatus::Changed,
707 info: Some(existing_account_updated_info.clone()),
708 previous_status: AccountStatus::Loaded,
709 previous_info: Some(existing_account_initial_info.clone()),
710 ..Default::default()
711 },
712 ),
713 (
714 existing_account_with_storage_address,
715 TransitionAccount {
716 status: AccountStatus::Changed,
717 info: Some(existing_account_with_storage_info.clone()),
718 previous_status: AccountStatus::Loaded,
719 previous_info: Some(existing_account_with_storage_info.clone()),
720 storage: HashMap::from_iter([
721 (
722 slot1,
723 StorageSlot::new_changed(StorageValue::from(1), StorageValue::from(10)),
724 ),
725 (
726 slot2,
727 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(20)),
728 ),
729 ]),
730 storage_was_destroyed: false,
731 },
732 ),
733 ]));
734
735 state.apply_transition(Vec::from([
737 (
738 new_account_address,
739 TransitionAccount {
740 status: AccountStatus::Destroyed,
741 info: None,
742 previous_status: AccountStatus::InMemoryChange,
743 previous_info: Some(new_account_created_info),
744 ..Default::default()
745 },
746 ),
747 (
748 existing_account_address,
749 TransitionAccount {
750 status: AccountStatus::Changed,
751 info: Some(existing_account_initial_info),
752 previous_status: AccountStatus::Changed,
753 previous_info: Some(existing_account_updated_info),
754 ..Default::default()
755 },
756 ),
757 (
758 existing_account_with_storage_address,
759 TransitionAccount {
760 status: AccountStatus::Changed,
761 info: Some(existing_account_with_storage_info.clone()),
762 previous_status: AccountStatus::Changed,
763 previous_info: Some(existing_account_with_storage_info.clone()),
764 storage: HashMap::from_iter([
765 (
766 slot1,
767 StorageSlot::new_changed(StorageValue::from(10), StorageValue::from(1)),
768 ),
769 (
770 slot2,
771 StorageSlot::new_changed(StorageValue::from(20), StorageValue::ZERO),
772 ),
773 ]),
774 storage_was_destroyed: false,
775 },
776 ),
777 ]));
778
779 state.merge_transitions(BundleRetention::Reverts);
780
781 let mut bundle_state = state.take_bundle();
782 bundle_state.reverts.sort();
783
784 assert_eq!(bundle_state.reverts.as_ref(), Vec::from([Vec::from([])]));
787 }
788
789 #[test]
791 fn selfdestruct_state_and_reverts() {
792 let mut state = State::builder().with_bundle_update().build();
793
794 let existing_account_address = Address::from_slice(&[0x1; 20]);
796 let existing_account_info = AccountInfo {
797 nonce: 1,
798 ..Default::default()
799 };
800
801 let (slot1, slot2) = (StorageKey::from(1), StorageKey::from(2));
802
803 state.apply_transition(Vec::from([(
805 existing_account_address,
806 TransitionAccount {
807 status: AccountStatus::Destroyed,
808 info: None,
809 previous_status: AccountStatus::Loaded,
810 previous_info: Some(existing_account_info.clone()),
811 storage: HashMap::default(),
812 storage_was_destroyed: true,
813 },
814 )]));
815
816 state.apply_transition(Vec::from([(
818 existing_account_address,
819 TransitionAccount {
820 status: AccountStatus::DestroyedChanged,
821 info: Some(existing_account_info.clone()),
822 previous_status: AccountStatus::Destroyed,
823 previous_info: None,
824 storage: HashMap::from_iter([(
825 slot1,
826 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(1)),
827 )]),
828 storage_was_destroyed: false,
829 },
830 )]));
831
832 state.apply_transition(Vec::from([(
834 existing_account_address,
835 TransitionAccount {
836 status: AccountStatus::DestroyedAgain,
837 info: None,
838 previous_status: AccountStatus::DestroyedChanged,
839 previous_info: Some(existing_account_info.clone()),
840 storage: HashMap::default(),
842 storage_was_destroyed: true,
843 },
844 )]));
845
846 state.apply_transition(Vec::from([(
848 existing_account_address,
849 TransitionAccount {
850 status: AccountStatus::DestroyedChanged,
851 info: Some(existing_account_info.clone()),
852 previous_status: AccountStatus::DestroyedAgain,
853 previous_info: None,
854 storage: HashMap::from_iter([(
855 slot2,
856 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(2)),
857 )]),
858 storage_was_destroyed: false,
859 },
860 )]));
861
862 state.merge_transitions(BundleRetention::Reverts);
863
864 let bundle_state = state.take_bundle();
865
866 assert_eq!(
867 bundle_state.state,
868 HashMap::from_iter([(
869 existing_account_address,
870 BundleAccount {
871 info: Some(existing_account_info.clone()),
872 original_info: Some(existing_account_info.clone()),
873 storage: HashMap::from_iter([(
874 slot2,
875 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(2))
876 )]),
877 status: AccountStatus::DestroyedChanged,
878 }
879 )])
880 );
881
882 assert_eq!(
883 bundle_state.reverts.as_ref(),
884 Vec::from([Vec::from([(
885 existing_account_address,
886 AccountRevert {
887 account: AccountInfoRevert::DoNothing,
888 previous_status: AccountStatus::Loaded,
889 storage: HashMap::from_iter([(slot2, RevertToSlot::Destroyed)]),
890 wipe_storage: true,
891 }
892 )])])
893 )
894 }
895}