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, 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) =
192 self.bundle_state.account(&address).cloned().map(Into::into)
193 {
194 return Ok(entry.insert(account));
195 }
196 }
197 let info = self.database.basic(address)?;
199 let account = match info {
200 None => CacheAccount::new_loaded_not_existing(),
201 Some(acc) if acc.is_empty() => {
202 CacheAccount::new_loaded_empty_eip161(HashMap::default())
203 }
204 Some(acc) => CacheAccount::new_loaded(acc, HashMap::default()),
205 };
206 Ok(entry.insert(account))
207 }
208 hash_map::Entry::Occupied(entry) => Ok(entry.into_mut()),
209 }
210 }
211
212 pub fn take_bundle(&mut self) -> BundleState {
224 core::mem::take(&mut self.bundle_state)
225 }
226}
227
228impl<DB: Database> Database for State<DB> {
229 type Error = DB::Error;
230
231 fn basic(&mut self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
232 self.load_cache_account(address).map(|a| a.account_info())
233 }
234
235 fn code_by_hash(&mut self, code_hash: B256) -> Result<Bytecode, Self::Error> {
236 let res = match self.cache.contracts.entry(code_hash) {
237 hash_map::Entry::Occupied(entry) => Ok(entry.get().clone()),
238 hash_map::Entry::Vacant(entry) => {
239 if self.use_preloaded_bundle {
240 if let Some(code) = self.bundle_state.contracts.get(&code_hash) {
241 entry.insert(code.clone());
242 return Ok(code.clone());
243 }
244 }
245 let code = self.database.code_by_hash(code_hash)?;
247 entry.insert(code.clone());
248 Ok(code)
249 }
250 };
251 res
252 }
253
254 fn storage(
255 &mut self,
256 address: Address,
257 index: StorageKey,
258 ) -> Result<StorageValue, Self::Error> {
259 if let Some(account) = self.cache.accounts.get_mut(&address) {
262 let is_storage_known = account.status.is_storage_known();
264 Ok(account
265 .account
266 .as_mut()
267 .map(|account| match account.storage.entry(index) {
268 hash_map::Entry::Occupied(entry) => Ok(*entry.get()),
269 hash_map::Entry::Vacant(entry) => {
270 let value = if is_storage_known {
273 StorageValue::ZERO
274 } else {
275 self.database.storage(address, index)?
276 };
277 entry.insert(value);
278 Ok(value)
279 }
280 })
281 .transpose()?
282 .unwrap_or_default())
283 } else {
284 unreachable!("For accessing any storage account is guaranteed to be loaded beforehand")
285 }
286 }
287
288 fn block_hash(&mut self, number: u64) -> Result<B256, Self::Error> {
289 match self.block_hashes.entry(number) {
290 btree_map::Entry::Occupied(entry) => Ok(*entry.get()),
291 btree_map::Entry::Vacant(entry) => {
292 let ret = *entry.insert(self.database.block_hash(number)?);
293
294 let last_block = number.saturating_sub(BLOCK_HASH_HISTORY);
296 while let Some(entry) = self.block_hashes.first_entry() {
297 if *entry.key() < last_block {
298 entry.remove();
299 } else {
300 break;
301 }
302 }
303
304 Ok(ret)
305 }
306 }
307 }
308}
309
310impl<DB: Database> DatabaseCommit for State<DB> {
311 fn commit(&mut self, evm_state: HashMap<Address, Account>) {
312 let transitions = self.cache.apply_evm_state(evm_state);
313 self.apply_transition(transitions);
314 }
315}
316
317#[cfg(test)]
318mod tests {
319 use super::*;
320 use crate::{
321 states::{reverts::AccountInfoRevert, StorageSlot},
322 AccountRevert, AccountStatus, BundleAccount, RevertToSlot,
323 };
324 use primitives::{keccak256, U256};
325
326 #[test]
327 fn block_hash_cache() {
328 let mut state = State::builder().build();
329 state.block_hash(1u64).unwrap();
330 state.block_hash(2u64).unwrap();
331
332 let test_number = BLOCK_HASH_HISTORY + 2;
333
334 let block1_hash = keccak256(U256::from(1).to_string().as_bytes());
335 let block2_hash = keccak256(U256::from(2).to_string().as_bytes());
336 let block_test_hash = keccak256(U256::from(test_number).to_string().as_bytes());
337
338 assert_eq!(
339 state.block_hashes,
340 BTreeMap::from([(1, block1_hash), (2, block2_hash)])
341 );
342
343 state.block_hash(test_number).unwrap();
344 assert_eq!(
345 state.block_hashes,
346 BTreeMap::from([(test_number, block_test_hash), (2, block2_hash)])
347 );
348 }
349
350 #[test]
357 fn reverts_preserve_old_values() {
358 let mut state = State::builder().with_bundle_update().build();
359
360 let (slot1, slot2, slot3) = (
361 StorageKey::from(1),
362 StorageKey::from(2),
363 StorageKey::from(3),
364 );
365
366 let new_account_address = Address::from_slice(&[0x1; 20]);
369 let new_account_created_info = AccountInfo {
370 nonce: 1,
371 balance: U256::from(1),
372 ..Default::default()
373 };
374 let new_account_changed_info = AccountInfo {
375 nonce: 2,
376 ..new_account_created_info.clone()
377 };
378 let new_account_changed_info2 = AccountInfo {
379 nonce: 3,
380 ..new_account_changed_info.clone()
381 };
382
383 let existing_account_address = Address::from_slice(&[0x2; 20]);
385 let existing_account_initial_info = AccountInfo {
386 nonce: 1,
387 ..Default::default()
388 };
389 let existing_account_initial_storage = HashMap::<StorageKey, StorageValue>::from_iter([
390 (slot1, StorageValue::from(100)), (slot2, StorageValue::from(200)), ]);
393 let existing_account_changed_info = AccountInfo {
394 nonce: 2,
395 ..existing_account_initial_info.clone()
396 };
397
398 state.apply_transition(Vec::from([
400 (
401 new_account_address,
402 TransitionAccount {
403 status: AccountStatus::InMemoryChange,
404 info: Some(new_account_created_info.clone()),
405 previous_status: AccountStatus::LoadedNotExisting,
406 previous_info: None,
407 ..Default::default()
408 },
409 ),
410 (
411 existing_account_address,
412 TransitionAccount {
413 status: AccountStatus::InMemoryChange,
414 info: Some(existing_account_changed_info.clone()),
415 previous_status: AccountStatus::Loaded,
416 previous_info: Some(existing_account_initial_info.clone()),
417 storage: HashMap::from_iter([(
418 slot1,
419 StorageSlot::new_changed(
420 *existing_account_initial_storage.get(&slot1).unwrap(),
421 StorageValue::from(1000),
422 ),
423 )]),
424 storage_was_destroyed: false,
425 },
426 ),
427 ]));
428
429 state.apply_transition(Vec::from([(
431 new_account_address,
432 TransitionAccount {
433 status: AccountStatus::InMemoryChange,
434 info: Some(new_account_changed_info.clone()),
435 previous_status: AccountStatus::InMemoryChange,
436 previous_info: Some(new_account_created_info.clone()),
437 ..Default::default()
438 },
439 )]));
440
441 state.apply_transition(Vec::from([
443 (
444 new_account_address,
445 TransitionAccount {
446 status: AccountStatus::InMemoryChange,
447 info: Some(new_account_changed_info2.clone()),
448 previous_status: AccountStatus::InMemoryChange,
449 previous_info: Some(new_account_changed_info),
450 storage: HashMap::from_iter([(
451 slot1,
452 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(1)),
453 )]),
454 storage_was_destroyed: false,
455 },
456 ),
457 (
458 existing_account_address,
459 TransitionAccount {
460 status: AccountStatus::InMemoryChange,
461 info: Some(existing_account_changed_info.clone()),
462 previous_status: AccountStatus::InMemoryChange,
463 previous_info: Some(existing_account_changed_info.clone()),
464 storage: HashMap::from_iter([
465 (
466 slot1,
467 StorageSlot::new_changed(
468 StorageValue::from(100),
469 StorageValue::from(1_000),
470 ),
471 ),
472 (
473 slot2,
474 StorageSlot::new_changed(
475 *existing_account_initial_storage.get(&slot2).unwrap(),
476 StorageValue::from(2_000),
477 ),
478 ),
479 (
481 slot3,
482 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(3_000)),
483 ),
484 ]),
485 storage_was_destroyed: false,
486 },
487 ),
488 ]));
489
490 state.merge_transitions(BundleRetention::Reverts);
491 let mut bundle_state = state.take_bundle();
492
493 bundle_state.reverts.sort();
496 assert_eq!(
497 bundle_state.reverts.as_ref(),
498 Vec::from([Vec::from([
499 (
500 new_account_address,
501 AccountRevert {
502 account: AccountInfoRevert::DeleteIt,
503 previous_status: AccountStatus::LoadedNotExisting,
504 storage: HashMap::from_iter([(
505 slot1,
506 RevertToSlot::Some(StorageValue::ZERO)
507 )]),
508 wipe_storage: false,
509 }
510 ),
511 (
512 existing_account_address,
513 AccountRevert {
514 account: AccountInfoRevert::RevertTo(existing_account_initial_info.clone()),
515 previous_status: AccountStatus::Loaded,
516 storage: HashMap::from_iter([
517 (
518 slot1,
519 RevertToSlot::Some(
520 *existing_account_initial_storage.get(&slot1).unwrap()
521 )
522 ),
523 (
524 slot2,
525 RevertToSlot::Some(
526 *existing_account_initial_storage.get(&slot2).unwrap()
527 )
528 ),
529 (slot3, RevertToSlot::Some(StorageValue::ZERO))
530 ]),
531 wipe_storage: false,
532 }
533 ),
534 ])]),
535 "The account or storage reverts are incorrect"
536 );
537
538 assert_eq!(
541 bundle_state.account(&new_account_address),
542 Some(&BundleAccount {
543 info: Some(new_account_changed_info2),
544 original_info: None,
545 status: AccountStatus::InMemoryChange,
546 storage: HashMap::from_iter([(
547 slot1,
548 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(1))
549 )]),
550 }),
551 "The latest state of the new account is incorrect"
552 );
553
554 assert_eq!(
557 bundle_state.account(&existing_account_address),
558 Some(&BundleAccount {
559 info: Some(existing_account_changed_info),
560 original_info: Some(existing_account_initial_info),
561 status: AccountStatus::InMemoryChange,
562 storage: HashMap::from_iter([
563 (
564 slot1,
565 StorageSlot::new_changed(
566 *existing_account_initial_storage.get(&slot1).unwrap(),
567 StorageValue::from(1_000)
568 )
569 ),
570 (
571 slot2,
572 StorageSlot::new_changed(
573 *existing_account_initial_storage.get(&slot2).unwrap(),
574 StorageValue::from(2_000)
575 )
576 ),
577 (
579 slot3,
580 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(3_000))
581 ),
582 ]),
583 }),
584 "The latest state of the existing account is incorrect"
585 );
586 }
587
588 #[test]
591 fn bundle_scoped_reverts_collapse() {
592 let mut state = State::builder().with_bundle_update().build();
593
594 let new_account_address = Address::from_slice(&[0x1; 20]);
596 let new_account_created_info = AccountInfo {
597 nonce: 1,
598 balance: U256::from(1),
599 ..Default::default()
600 };
601
602 let existing_account_address = Address::from_slice(&[0x2; 20]);
604 let existing_account_initial_info = AccountInfo {
605 nonce: 1,
606 ..Default::default()
607 };
608 let existing_account_updated_info = AccountInfo {
609 nonce: 1,
610 balance: U256::from(1),
611 ..Default::default()
612 };
613
614 let (slot1, slot2) = (StorageKey::from(1), StorageKey::from(2));
616 let existing_account_with_storage_address = Address::from_slice(&[0x3; 20]);
617 let existing_account_with_storage_info = AccountInfo {
618 nonce: 1,
619 ..Default::default()
620 };
621 state.apply_transition(Vec::from([
623 (
624 new_account_address,
625 TransitionAccount {
626 status: AccountStatus::InMemoryChange,
627 info: Some(new_account_created_info.clone()),
628 previous_status: AccountStatus::LoadedNotExisting,
629 previous_info: None,
630 ..Default::default()
631 },
632 ),
633 (
634 existing_account_address,
635 TransitionAccount {
636 status: AccountStatus::Changed,
637 info: Some(existing_account_updated_info.clone()),
638 previous_status: AccountStatus::Loaded,
639 previous_info: Some(existing_account_initial_info.clone()),
640 ..Default::default()
641 },
642 ),
643 (
644 existing_account_with_storage_address,
645 TransitionAccount {
646 status: AccountStatus::Changed,
647 info: Some(existing_account_with_storage_info.clone()),
648 previous_status: AccountStatus::Loaded,
649 previous_info: Some(existing_account_with_storage_info.clone()),
650 storage: HashMap::from_iter([
651 (
652 slot1,
653 StorageSlot::new_changed(StorageValue::from(1), StorageValue::from(10)),
654 ),
655 (
656 slot2,
657 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(20)),
658 ),
659 ]),
660 storage_was_destroyed: false,
661 },
662 ),
663 ]));
664
665 state.apply_transition(Vec::from([
667 (
668 new_account_address,
669 TransitionAccount {
670 status: AccountStatus::Destroyed,
671 info: None,
672 previous_status: AccountStatus::InMemoryChange,
673 previous_info: Some(new_account_created_info),
674 ..Default::default()
675 },
676 ),
677 (
678 existing_account_address,
679 TransitionAccount {
680 status: AccountStatus::Changed,
681 info: Some(existing_account_initial_info),
682 previous_status: AccountStatus::Changed,
683 previous_info: Some(existing_account_updated_info),
684 ..Default::default()
685 },
686 ),
687 (
688 existing_account_with_storage_address,
689 TransitionAccount {
690 status: AccountStatus::Changed,
691 info: Some(existing_account_with_storage_info.clone()),
692 previous_status: AccountStatus::Changed,
693 previous_info: Some(existing_account_with_storage_info.clone()),
694 storage: HashMap::from_iter([
695 (
696 slot1,
697 StorageSlot::new_changed(StorageValue::from(10), StorageValue::from(1)),
698 ),
699 (
700 slot2,
701 StorageSlot::new_changed(StorageValue::from(20), StorageValue::ZERO),
702 ),
703 ]),
704 storage_was_destroyed: false,
705 },
706 ),
707 ]));
708
709 state.merge_transitions(BundleRetention::Reverts);
710
711 let mut bundle_state = state.take_bundle();
712 bundle_state.reverts.sort();
713
714 assert_eq!(bundle_state.reverts.as_ref(), Vec::from([Vec::from([])]));
717 }
718
719 #[test]
721 fn selfdestruct_state_and_reverts() {
722 let mut state = State::builder().with_bundle_update().build();
723
724 let existing_account_address = Address::from_slice(&[0x1; 20]);
726 let existing_account_info = AccountInfo {
727 nonce: 1,
728 ..Default::default()
729 };
730
731 let (slot1, slot2) = (StorageKey::from(1), StorageKey::from(2));
732
733 state.apply_transition(Vec::from([(
735 existing_account_address,
736 TransitionAccount {
737 status: AccountStatus::Destroyed,
738 info: None,
739 previous_status: AccountStatus::Loaded,
740 previous_info: Some(existing_account_info.clone()),
741 storage: HashMap::default(),
742 storage_was_destroyed: true,
743 },
744 )]));
745
746 state.apply_transition(Vec::from([(
748 existing_account_address,
749 TransitionAccount {
750 status: AccountStatus::DestroyedChanged,
751 info: Some(existing_account_info.clone()),
752 previous_status: AccountStatus::Destroyed,
753 previous_info: None,
754 storage: HashMap::from_iter([(
755 slot1,
756 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(1)),
757 )]),
758 storage_was_destroyed: false,
759 },
760 )]));
761
762 state.apply_transition(Vec::from([(
764 existing_account_address,
765 TransitionAccount {
766 status: AccountStatus::DestroyedAgain,
767 info: None,
768 previous_status: AccountStatus::DestroyedChanged,
769 previous_info: Some(existing_account_info.clone()),
770 storage: HashMap::default(),
772 storage_was_destroyed: true,
773 },
774 )]));
775
776 state.apply_transition(Vec::from([(
778 existing_account_address,
779 TransitionAccount {
780 status: AccountStatus::DestroyedChanged,
781 info: Some(existing_account_info.clone()),
782 previous_status: AccountStatus::DestroyedAgain,
783 previous_info: None,
784 storage: HashMap::from_iter([(
785 slot2,
786 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(2)),
787 )]),
788 storage_was_destroyed: false,
789 },
790 )]));
791
792 state.merge_transitions(BundleRetention::Reverts);
793
794 let bundle_state = state.take_bundle();
795
796 assert_eq!(
797 bundle_state.state,
798 HashMap::from_iter([(
799 existing_account_address,
800 BundleAccount {
801 info: Some(existing_account_info.clone()),
802 original_info: Some(existing_account_info.clone()),
803 storage: HashMap::from_iter([(
804 slot2,
805 StorageSlot::new_changed(StorageValue::ZERO, StorageValue::from(2))
806 )]),
807 status: AccountStatus::DestroyedChanged,
808 }
809 )])
810 );
811
812 assert_eq!(
813 bundle_state.reverts.as_ref(),
814 Vec::from([Vec::from([(
815 existing_account_address,
816 AccountRevert {
817 account: AccountInfoRevert::DoNothing,
818 previous_status: AccountStatus::Loaded,
819 storage: HashMap::from_iter([(slot2, RevertToSlot::Destroyed)]),
820 wipe_storage: true,
821 }
822 )])])
823 )
824 }
825}