ink_engine/
database.rs

1// Copyright (C) Use Ink (UK) Ltd.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::types::Balance;
16use ink_primitives::{
17    AccountId,
18    Address,
19    H256,
20    U256,
21};
22use scale::KeyedVec;
23use std::collections::HashMap;
24
25const BALANCE_OF: &[u8] = b"balance:";
26const STORAGE_OF: &[u8] = b"contract-storage:";
27const CONTRACT_PREFIX: &[u8] = b"contract:";
28const MSG_HANDLER_OF: &[u8] = b"message-handler:";
29const CODE_HASH_OF: &[u8] = b"code-hash:";
30
31/// Returns the database key under which to find the balance for contract `who`.
32pub fn balance_of_key(who: &Address) -> [u8; 32] {
33    let keyed = who.0.to_vec().to_keyed_vec(BALANCE_OF);
34    let mut hashed_key: [u8; 32] = [0; 32];
35    super::hashing::blake2b_256(&keyed[..], &mut hashed_key);
36    hashed_key
37}
38
39/// Returns the database key under which to find the storage for contract `who`.
40pub fn storage_of_contract_key(who: &Address, key: &[u8]) -> [u8; 32] {
41    let keyed = who
42        .as_bytes()
43        .to_vec()
44        .to_keyed_vec(key)
45        .to_keyed_vec(STORAGE_OF);
46    let mut hashed_key: [u8; 32] = [0; 32];
47    super::hashing::blake2b_256(&keyed[..], &mut hashed_key);
48    hashed_key
49}
50
51pub type MessageHandler = fn(Vec<u8>) -> Vec<u8>;
52
53pub fn contract_key(f: MessageHandler) -> [u8; 32] {
54    let f = f as usize;
55    let f = f.to_le_bytes();
56    let keyed = f.to_vec().to_keyed_vec(CONTRACT_PREFIX);
57    let mut ret: [u8; 32] = [0; 32];
58    super::hashing::blake2b_256(&keyed[..], &mut ret);
59    ret
60}
61
62pub fn message_handler_of_contract_key(key: &[u8]) -> [u8; 32] {
63    let keyed = key.to_vec().to_keyed_vec(MSG_HANDLER_OF);
64    let mut hashed_key: [u8; 32] = [0; 32];
65    super::hashing::blake2b_256(&keyed[..], &mut hashed_key);
66    hashed_key
67}
68
69pub fn code_hash_for_addr(addr: &Address) -> [u8; 32] {
70    let key = addr.0;
71    let keyed = key.to_keyed_vec(CODE_HASH_OF);
72    let mut hashed_key: [u8; 32] = [0; 32];
73    super::hashing::blake2b_256(&keyed[..], &mut hashed_key);
74    hashed_key
75}
76
77/// The chain database.
78///
79/// Everything is stored in here: contracts, balances, contract storage, etc.
80/// Just like in Substrate a prefix hash is computed for every contract.
81#[derive(Default)]
82pub struct Database {
83    hmap: HashMap<Vec<u8>, Vec<u8>>,
84    fmap: HashMap<Vec<u8>, MessageHandler>,
85}
86
87impl Database {
88    /// Creates a new database instance.
89    pub fn new() -> Self {
90        Database {
91            hmap: HashMap::new(),
92            fmap: HashMap::new(),
93        }
94    }
95
96    /// Returns the amount of entries in the database.
97    #[cfg(test)]
98    fn len(&self) -> usize {
99        self.hmap.len()
100    }
101
102    /// Returns a reference to the value corresponding to the key.
103    fn get(&self, key: &[u8]) -> Option<&Vec<u8>> {
104        self.hmap.get(key)
105    }
106
107    /// Returns a reference to the value corresponding to the key.
108    pub fn get_from_contract_storage(
109        &self,
110        addr: &Address,
111        key: &[u8],
112    ) -> Option<&Vec<u8>> {
113        let hashed_key = storage_of_contract_key(addr, key);
114        self.hmap.get(hashed_key.as_slice())
115    }
116
117    /// Inserts `value` into the contract storage of `addr` at storage key `key`.
118    pub fn insert_into_contract_storage(
119        &mut self,
120        addr: &Address,
121        key: &[u8],
122        value: Vec<u8>,
123    ) -> Option<Vec<u8>> {
124        let hashed_key = storage_of_contract_key(addr, key);
125        self.hmap.insert(hashed_key.to_vec(), value)
126    }
127
128    /// Removes the value at the contract storage of `addr` at storage key `key`.
129    pub fn remove_contract_storage(
130        &mut self,
131        addr: &Address,
132        key: &[u8],
133    ) -> Option<Vec<u8>> {
134        let hashed_key = storage_of_contract_key(addr, key);
135        self.hmap.remove(hashed_key.as_slice())
136    }
137
138    /// Removes a key from the storage, returning the value at the key if the key
139    /// was previously in storage.
140    pub fn remove(&mut self, key: &[u8]) -> Option<Vec<u8>> {
141        self.hmap.remove(key)
142    }
143
144    /// Sets the value of the entry, and returns the entry's old value.
145    pub fn insert(&mut self, key: Vec<u8>, value: Vec<u8>) -> Option<Vec<u8>> {
146        self.hmap.insert(key, value)
147    }
148
149    /// Clears the database, removing all key-value pairs.
150    pub fn clear(&mut self) {
151        self.hmap.clear();
152    }
153
154    /// Returns the balance of the contract at `addr`, if available.
155    pub fn get_acc_balance(&self, _addr: &AccountId) -> Option<Balance> {
156        todo!()
157    }
158
159    /// Sets the balance of `addr` to `new_balance`.
160    pub fn set_acc_balance(&mut self, _addr: &AccountId, _new_balance: Balance) {
161        todo!()
162    }
163
164    pub fn get_balance(&self, addr: &Address) -> Option<U256> {
165        let hashed_key = balance_of_key(addr);
166        self.get(&hashed_key).map(|encoded_balance| {
167            scale::Decode::decode(&mut &encoded_balance[..])
168                .expect("unable to decode balance from database")
169        })
170    }
171
172    /// Sets the balance of `addr` to `new_balance`.
173    pub fn set_balance(&mut self, addr: &Address, new_balance: U256) {
174        let hashed_key = balance_of_key(addr);
175        let encoded_balance = scale::Encode::encode(&new_balance);
176        self.hmap
177            .entry(hashed_key.to_vec())
178            .and_modify(|v| *v = encoded_balance.clone())
179            .or_insert(encoded_balance);
180    }
181
182    pub fn set_contract_message_handler(&mut self, handler: MessageHandler) -> [u8; 32] {
183        let key = contract_key(handler);
184        let hashed_key = message_handler_of_contract_key(&key);
185        self.fmap
186            .entry(hashed_key.to_vec())
187            .and_modify(|x| *x = handler)
188            .or_insert(handler);
189        key
190    }
191
192    /// Returns the message handler for a code hash.
193    pub fn get_contract_message_handler(&mut self, code_hash: &H256) -> MessageHandler {
194        let hashed_key = message_handler_of_contract_key(&code_hash.0);
195        *self.fmap.get(hashed_key.as_slice()).unwrap()
196    }
197
198    pub fn set_code_hash(&mut self, addr: &Address, code_hash: &H256) {
199        let hashed_key = code_hash_for_addr(addr);
200        self.hmap
201            .entry(hashed_key.to_vec())
202            .and_modify(|x| *x = code_hash.as_bytes().to_vec())
203            .or_insert(code_hash.as_bytes().to_vec());
204    }
205
206    pub fn get_code_hash(&self, addr: &Address) -> Option<H256> {
207        let hashed_key = code_hash_for_addr(addr);
208        self.get(&hashed_key)
209            .cloned()
210            .map(|v| H256::from_slice(v.as_slice()))
211    }
212}
213
214#[cfg(test)]
215mod tests {
216    use super::{
217        Address,
218        Database,
219    };
220
221    #[test]
222    fn basic_operations() {
223        let mut database = Database::new();
224        let key1 = vec![42];
225        let key2 = vec![43];
226        let val1 = vec![44];
227        let val2 = vec![45];
228        let val3 = vec![46];
229
230        assert_eq!(database.len(), 0);
231        assert_eq!(database.get(&key1), None);
232        assert_eq!(database.insert(key1.clone(), val1.clone()), None);
233        assert_eq!(database.get(&key1), Some(&val1));
234        assert_eq!(database.insert(key1.clone(), val2.clone()), Some(val1));
235        assert_eq!(database.get(&key1), Some(&val2));
236        assert_eq!(database.insert(key2.clone(), val3.clone()), None);
237        assert_eq!(database.len(), 2);
238        assert_eq!(database.remove(&key2), Some(val3));
239        assert_eq!(database.len(), 1);
240        database.clear();
241        assert_eq!(database.len(), 0);
242    }
243
244    #[test]
245    fn contract_storage() {
246        let addr = Address::from([1; 20]);
247        let mut storage = Database::new();
248        let key1 = vec![42];
249        let key2 = vec![43];
250        let val1 = vec![44];
251        let val2 = vec![45];
252        let val3 = vec![46];
253
254        assert_eq!(storage.len(), 0);
255        assert_eq!(storage.get_from_contract_storage(&addr, &key1), None);
256        assert_eq!(
257            storage.insert_into_contract_storage(&addr, &key1, val1.clone()),
258            None
259        );
260        assert_eq!(storage.get_from_contract_storage(&addr, &key1), Some(&val1));
261        assert_eq!(
262            storage.insert_into_contract_storage(&addr, &key1, val2.clone()),
263            Some(val1)
264        );
265        assert_eq!(storage.get_from_contract_storage(&addr, &key1), Some(&val2));
266        assert_eq!(
267            storage.insert_into_contract_storage(&addr, &key2, val3.clone()),
268            None
269        );
270        assert_eq!(storage.len(), 2);
271        assert_eq!(storage.remove_contract_storage(&addr, &key2), Some(val3));
272        assert_eq!(storage.len(), 1);
273        assert_eq!(storage.remove_contract_storage(&addr, &key1), Some(val2));
274        assert_eq!(storage.len(), 0);
275    }
276}