1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153

use crate::{subroutine::Filth, Database, KECCAK_EMPTY};

use alloc::vec::Vec;
use hashbrown::{hash_map::Entry, HashMap as Map};

use primitive_types::{H160, H256, U256};

use crate::{Account, AccountInfo, Log};
use bytes::Bytes;
use sha3::{Digest, Keccak256};

use super::DatabaseCommit;

/// Memory backend, storing all state values in a `Map` in memory.
#[derive(Debug, Clone)]
pub struct InMemoryDB {
    /// dummy account info where code is allways None. Code bytes can be found in `contracts`
    cache: Map<H160, AccountInfo>,
    storage: Map<H160, Map<U256, U256>>,
    contracts: Map<H256, Bytes>,
    logs: Vec<Log>,
}

impl Default for InMemoryDB {
    fn default() -> Self {
        Self {
            cache: Map::new(),
            storage: Map::new(),
            contracts: Map::new(),
            logs: Vec::default(),
        }
    }
}

impl InMemoryDB {
    pub fn cache(&self) -> &Map<H160, AccountInfo> {
        &self.cache
    }
    pub fn storage(&self) -> &Map<H160, Map<U256, U256>> {
        &self.storage
    }

    pub fn insert_cache(&mut self, address: H160, mut account: AccountInfo) {
        let code = core::mem::take(&mut account.code);
        if let Some(code) = code {
            if !code.is_empty() {
                let code_hash = H256::from_slice(&Keccak256::digest(&code));
                account.code_hash = code_hash;
                self.contracts.insert(code_hash, code);
            }
        }
        if account.code_hash.is_zero() {
            account.code_hash = KECCAK_EMPTY;
        }
        self.cache.insert(address, account);
    }

    pub fn insert_cache_storage(&mut self, address: H160, slot: U256, value: U256) {
        self.storage.entry(address).or_default().insert(slot, value);
    }

    /// Create a new memory backend.
    pub fn new() -> Self {
        let mut contracts = Map::new();
        contracts.insert(KECCAK_EMPTY, Bytes::new());
        contracts.insert(H256::zero(), Bytes::new());
        Self {
            contracts,
            ..core::default::Default::default()
        }
    }

    /// return true if account exists or fetch it from database
    fn fetch_account(&mut self, address: &H160) -> bool {
        {
            if let Some(acc) = self.cache.get(address) {
                return acc.exists();
            }
        }
        false
    }
}

impl DatabaseCommit for InMemoryDB {
    fn commit(&mut self, changes: Map<H160, Account>) {
        for (add, acc) in changes {
            if acc.is_empty() || matches!(acc.filth, Filth::Destroyed) {
                self.cache.remove(&add);
                self.storage.remove(&add);
            } else {
                self.insert_cache(add, acc.info);
                let storage = self.storage.entry(add).or_default();
                if acc.filth.abandon_old_storage() {
                    storage.clear();
                }
                for (index, value) in acc.storage {
                    if value.is_zero() {
                        storage.remove(&index);
                    } else {
                        storage.insert(index, value);
                    }
                }
                if storage.is_empty() {
                    self.storage.remove(&add);
                }
            }
        }
    }
}

impl Database for InMemoryDB {
    fn block_hash(&mut self, _number: U256) -> H256 {
        H256::zero()
    }

    fn exists(&mut self, address: H160) -> Option<AccountInfo> {
        if self.fetch_account(&address) {
            Some(self.cache.get(&address).cloned().unwrap())
        } else {
            None
        }
    }

    fn basic(&mut self, address: H160) -> AccountInfo {
        if self.fetch_account(&address) {
            let mut basic = self.cache.get(&address).cloned().unwrap();
            basic.code = None;
            basic
        } else {
            AccountInfo::default()
        }
    }

    fn storage(&mut self, address: H160, index: U256) -> U256 {
        if self.fetch_account(&address) {
            if let Some(storage) = self.storage.get(&address) {
                if let Some(slot) = storage.get(&index) {
                    return *slot;
                }
            }
            U256::zero()
        } else {
            U256::zero()
        }
    }

    fn code_by_hash(&mut self, code_hash: H256) -> Bytes {
        match self.contracts.entry(code_hash) {
            Entry::Occupied(entry) => entry.get().clone(),
            Entry::Vacant(_entry) => Bytes::new(),
        }
    }
}