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:";
30
31pub 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
39pub 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#[derive(Default)]
82pub struct Database {
83 hmap: HashMap<Vec<u8>, Vec<u8>>,
84 fmap: HashMap<Vec<u8>, MessageHandler>,
85}
86
87impl Database {
88 pub fn new() -> Self {
90 Database {
91 hmap: HashMap::new(),
92 fmap: HashMap::new(),
93 }
94 }
95
96 #[cfg(test)]
98 fn len(&self) -> usize {
99 self.hmap.len()
100 }
101
102 fn get(&self, key: &[u8]) -> Option<&Vec<u8>> {
104 self.hmap.get(key)
105 }
106
107 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 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 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 pub fn remove(&mut self, key: &[u8]) -> Option<Vec<u8>> {
141 self.hmap.remove(key)
142 }
143
144 pub fn insert(&mut self, key: Vec<u8>, value: Vec<u8>) -> Option<Vec<u8>> {
146 self.hmap.insert(key, value)
147 }
148
149 pub fn clear(&mut self) {
151 self.hmap.clear();
152 }
153
154 pub fn get_acc_balance(&self, _addr: &AccountId) -> Option<Balance> {
156 todo!()
157 }
158
159 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 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 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}