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:";
30const ACCOUNT_ID_OF: &[u8] = b"account-id:";
31
32/// Returns the database key under which to find the balance for contract `who`.
33pub fn balance_of_key(who: &Address) -> [u8; 32] {
34    let keyed = who.0.to_vec().to_keyed_vec(BALANCE_OF);
35    let mut hashed_key: [u8; 32] = [0; 32];
36    super::hashing::blake2b_256(&keyed[..], &mut hashed_key);
37    hashed_key
38}
39
40/// Returns the database key under which to find the account id for the address `who`.
41pub fn account_id_of_key(who: &Address) -> [u8; 32] {
42    let keyed = who.0.to_vec().to_keyed_vec(ACCOUNT_ID_OF);
43    let mut hashed_key: [u8; 32] = [0; 32];
44    super::hashing::blake2b_256(&keyed[..], &mut hashed_key);
45    hashed_key
46}
47
48/// Returns the database key under which to find the storage for contract `who`.
49pub fn storage_of_contract_key(who: &Address, key: &[u8]) -> [u8; 32] {
50    let keyed = who
51        .as_bytes()
52        .to_vec()
53        .to_keyed_vec(key)
54        .to_keyed_vec(STORAGE_OF);
55    let mut hashed_key: [u8; 32] = [0; 32];
56    super::hashing::blake2b_256(&keyed[..], &mut hashed_key);
57    hashed_key
58}
59
60pub type MessageHandler = fn(Vec<u8>) -> Vec<u8>;
61
62pub fn contract_key(f: MessageHandler) -> [u8; 32] {
63    let f = f as usize;
64    let f = f.to_le_bytes();
65    let keyed = f.to_vec().to_keyed_vec(CONTRACT_PREFIX);
66    let mut ret: [u8; 32] = [0; 32];
67    super::hashing::blake2b_256(&keyed[..], &mut ret);
68    ret
69}
70
71pub fn message_handler_of_contract_key(key: &[u8]) -> [u8; 32] {
72    let keyed = key.to_vec().to_keyed_vec(MSG_HANDLER_OF);
73    let mut hashed_key: [u8; 32] = [0; 32];
74    super::hashing::blake2b_256(&keyed[..], &mut hashed_key);
75    hashed_key
76}
77
78pub fn code_hash_for_addr(addr: &Address) -> [u8; 32] {
79    let key = addr.0;
80    let keyed = key.to_keyed_vec(CODE_HASH_OF);
81    let mut hashed_key: [u8; 32] = [0; 32];
82    super::hashing::blake2b_256(&keyed[..], &mut hashed_key);
83    hashed_key
84}
85
86/// The chain database.
87///
88/// Everything is stored in here: contracts, balances, contract storage, etc.
89/// Just like in Substrate a prefix hash is computed for every contract.
90#[derive(Default)]
91pub struct Database {
92    hmap: HashMap<Vec<u8>, Vec<u8>>,
93    fmap: HashMap<Vec<u8>, MessageHandler>,
94}
95
96impl Database {
97    /// Creates a new database instance.
98    pub fn new() -> Self {
99        Database {
100            hmap: HashMap::new(),
101            fmap: HashMap::new(),
102        }
103    }
104
105    /// Returns the amount of entries in the database.
106    #[cfg(test)]
107    fn len(&self) -> usize {
108        self.hmap.len()
109    }
110
111    /// Returns a reference to the value corresponding to the key.
112    fn get(&self, key: &[u8]) -> Option<&Vec<u8>> {
113        self.hmap.get(key)
114    }
115
116    /// Returns a reference to the value corresponding to the key.
117    pub fn get_from_contract_storage(
118        &self,
119        addr: &Address,
120        key: &[u8],
121    ) -> Option<&Vec<u8>> {
122        let hashed_key = storage_of_contract_key(addr, key);
123        self.hmap.get(hashed_key.as_slice())
124    }
125
126    /// Inserts `value` into the contract storage of `addr` at storage key `key`.
127    pub fn insert_into_contract_storage(
128        &mut self,
129        addr: &Address,
130        key: &[u8],
131        value: Vec<u8>,
132    ) -> Option<Vec<u8>> {
133        let hashed_key = storage_of_contract_key(addr, key);
134        self.hmap.insert(hashed_key.to_vec(), value)
135    }
136
137    /// Removes the value at the contract storage of `addr` at storage key `key`.
138    pub fn remove_contract_storage(
139        &mut self,
140        addr: &Address,
141        key: &[u8],
142    ) -> Option<Vec<u8>> {
143        let hashed_key = storage_of_contract_key(addr, key);
144        self.hmap.remove(hashed_key.as_slice())
145    }
146
147    /// Removes a key from the storage, returning the value at the key if the key
148    /// was previously in storage.
149    pub fn remove(&mut self, key: &[u8]) -> Option<Vec<u8>> {
150        self.hmap.remove(key)
151    }
152
153    /// Sets the value of the entry, and returns the entry's old value.
154    pub fn insert(&mut self, key: Vec<u8>, value: Vec<u8>) -> Option<Vec<u8>> {
155        self.hmap.insert(key, value)
156    }
157
158    /// Clears the database, removing all key-value pairs.
159    pub fn clear(&mut self) {
160        self.hmap.clear();
161    }
162
163    /// Returns the balance of the contract at `addr`, if available.
164    pub fn get_acc_balance(&self, _account_id: &AccountId) -> Option<Balance> {
165        todo!()
166    }
167
168    /// Sets the balance of `addr` to `new_balance`.
169    pub fn set_acc_balance(&mut self, _account_id: &AccountId, _new_balance: Balance) {
170        todo!()
171    }
172
173    /// Retrieves the account id for a specified address.
174    pub fn to_account_id(&self, addr: &Address) -> Vec<u8> {
175        let hashed_key = account_id_of_key(addr);
176        self.get(&hashed_key).cloned().unwrap_or_else(|| {
177            let mut bytes = [0xEE; 32];
178            bytes[..20].copy_from_slice(&addr.as_bytes()[..20]);
179            Vec::from(bytes)
180        })
181    }
182
183    pub fn get_balance(&self, addr: &Address) -> Option<U256> {
184        let hashed_key = balance_of_key(addr);
185        self.get(&hashed_key).map(|encoded_balance| {
186            scale::Decode::decode(&mut &encoded_balance[..])
187                .expect("unable to decode balance from database")
188        })
189    }
190
191    /// Sets the balance of `addr` to `new_balance`.
192    pub fn set_balance(&mut self, addr: &Address, new_balance: U256) {
193        let hashed_key = balance_of_key(addr);
194        let encoded_balance = scale::Encode::encode(&new_balance);
195        self.hmap
196            .entry(hashed_key.to_vec())
197            .and_modify(|v| *v = encoded_balance.clone())
198            .or_insert(encoded_balance);
199    }
200
201    pub fn set_contract_message_handler(&mut self, handler: MessageHandler) -> [u8; 32] {
202        let key = contract_key(handler);
203        let hashed_key = message_handler_of_contract_key(&key);
204        self.fmap
205            .entry(hashed_key.to_vec())
206            .and_modify(|x| *x = handler)
207            .or_insert(handler);
208        key
209    }
210
211    /// Returns the message handler for a code hash.
212    pub fn get_contract_message_handler(&mut self, code_hash: &H256) -> MessageHandler {
213        let hashed_key = message_handler_of_contract_key(&code_hash.0);
214        *self.fmap.get(hashed_key.as_slice()).unwrap()
215    }
216
217    pub fn set_code_hash(&mut self, addr: &Address, code_hash: &H256) {
218        let hashed_key = code_hash_for_addr(addr);
219        self.hmap
220            .entry(hashed_key.to_vec())
221            .and_modify(|x| *x = code_hash.as_bytes().to_vec())
222            .or_insert(code_hash.as_bytes().to_vec());
223    }
224
225    pub fn get_code_hash(&self, addr: &Address) -> Option<H256> {
226        let hashed_key = code_hash_for_addr(addr);
227        self.get(&hashed_key)
228            .cloned()
229            .map(|v| H256::from_slice(v.as_slice()))
230    }
231}
232
233#[cfg(test)]
234mod tests {
235    use super::{
236        Address,
237        Database,
238    };
239
240    #[test]
241    fn basic_operations() {
242        let mut database = Database::new();
243        let key1 = vec![42];
244        let key2 = vec![43];
245        let val1 = vec![44];
246        let val2 = vec![45];
247        let val3 = vec![46];
248
249        assert_eq!(database.len(), 0);
250        assert_eq!(database.get(&key1), None);
251        assert_eq!(database.insert(key1.clone(), val1.clone()), None);
252        assert_eq!(database.get(&key1), Some(&val1));
253        assert_eq!(database.insert(key1.clone(), val2.clone()), Some(val1));
254        assert_eq!(database.get(&key1), Some(&val2));
255        assert_eq!(database.insert(key2.clone(), val3.clone()), None);
256        assert_eq!(database.len(), 2);
257        assert_eq!(database.remove(&key2), Some(val3));
258        assert_eq!(database.len(), 1);
259        database.clear();
260        assert_eq!(database.len(), 0);
261    }
262
263    #[test]
264    fn contract_storage() {
265        let addr = Address::from([1; 20]);
266        let mut storage = Database::new();
267        let key1 = vec![42];
268        let key2 = vec![43];
269        let val1 = vec![44];
270        let val2 = vec![45];
271        let val3 = vec![46];
272
273        assert_eq!(storage.len(), 0);
274        assert_eq!(storage.get_from_contract_storage(&addr, &key1), None);
275        assert_eq!(
276            storage.insert_into_contract_storage(&addr, &key1, val1.clone()),
277            None
278        );
279        assert_eq!(storage.get_from_contract_storage(&addr, &key1), Some(&val1));
280        assert_eq!(
281            storage.insert_into_contract_storage(&addr, &key1, val2.clone()),
282            Some(val1)
283        );
284        assert_eq!(storage.get_from_contract_storage(&addr, &key1), Some(&val2));
285        assert_eq!(
286            storage.insert_into_contract_storage(&addr, &key2, val3.clone()),
287            None
288        );
289        assert_eq!(storage.len(), 2);
290        assert_eq!(storage.remove_contract_storage(&addr, &key2), Some(val3));
291        assert_eq!(storage.len(), 1);
292        assert_eq!(storage.remove_contract_storage(&addr, &key1), Some(val2));
293        assert_eq!(storage.len(), 0);
294    }
295}