1use 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
32pub 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
40pub 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
48pub 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#[derive(Default)]
91pub struct Database {
92 hmap: HashMap<Vec<u8>, Vec<u8>>,
93 fmap: HashMap<Vec<u8>, MessageHandler>,
94}
95
96impl Database {
97 pub fn new() -> Self {
99 Database {
100 hmap: HashMap::new(),
101 fmap: HashMap::new(),
102 }
103 }
104
105 #[cfg(test)]
107 fn len(&self) -> usize {
108 self.hmap.len()
109 }
110
111 fn get(&self, key: &[u8]) -> Option<&Vec<u8>> {
113 self.hmap.get(key)
114 }
115
116 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 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 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 pub fn remove(&mut self, key: &[u8]) -> Option<Vec<u8>> {
150 self.hmap.remove(key)
151 }
152
153 pub fn insert(&mut self, key: Vec<u8>, value: Vec<u8>) -> Option<Vec<u8>> {
155 self.hmap.insert(key, value)
156 }
157
158 pub fn clear(&mut self) {
160 self.hmap.clear();
161 }
162
163 pub fn get_acc_balance(&self, _account_id: &AccountId) -> Option<Balance> {
165 todo!()
166 }
167
168 pub fn set_acc_balance(&mut self, _account_id: &AccountId, _new_balance: Balance) {
170 todo!()
171 }
172
173 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 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 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}