1use std::{
21 collections::BTreeMap, any::{TypeId, Any}, iter::FromIterator, ops::Bound,
22};
23use crate::{Backend, StorageKey, StorageValue};
24use tetsy_hash_db::Hasher;
25use tp_trie::{TrieConfiguration, empty_child_trie_root};
26use tp_trie::trie_types::Layout;
27use tet_core::{
28 storage::{
29 well_known_keys::is_child_storage_key, Storage,
30 ChildInfo, StorageChild, TrackedStorageKey,
31 },
32 traits::Externalities, Blake2Hasher,
33};
34use log::warn;
35use codec::Encode;
36use externalities::{Extensions, Extension};
37
38#[derive(Debug)]
40pub struct BasicExternalities {
41 inner: Storage,
42 extensions: Extensions,
43}
44
45impl BasicExternalities {
46 pub fn new(inner: Storage) -> Self {
48 BasicExternalities { inner, extensions: Default::default() }
49 }
50
51 pub fn new_empty() -> Self {
53 Self::new(Storage::default())
54 }
55
56 pub fn insert(&mut self, k: StorageKey, v: StorageValue) -> Option<StorageValue> {
58 self.inner.top.insert(k, v)
59 }
60
61 pub fn into_storages(self) -> Storage {
63 self.inner
64 }
65
66 pub fn execute_with_storage<R>(
70 storage: &mut tet_core::storage::Storage,
71 f: impl FnOnce() -> R,
72 ) -> R {
73 let mut ext = Self {
74 inner: Storage {
75 top: std::mem::take(&mut storage.top),
76 children_default: std::mem::take(&mut storage.children_default),
77 },
78 extensions: Default::default(),
79 };
80
81 let r = ext.execute_with(f);
82
83 *storage = ext.into_storages();
84
85 r
86 }
87
88 pub fn execute_with<R>(&mut self, f: impl FnOnce() -> R) -> R {
92 externalities::set_and_run_with_externalities(self, f)
93 }
94
95 pub fn extensions(&mut self) -> &mut Extensions {
97 &mut self.extensions
98 }
99
100 pub fn register_extension(&mut self, ext: impl Extension) {
102 self.extensions.register(ext);
103 }
104}
105
106impl PartialEq for BasicExternalities {
107 fn eq(&self, other: &BasicExternalities) -> bool {
108 self.inner.top.eq(&other.inner.top)
109 && self.inner.children_default.eq(&other.inner.children_default)
110 }
111}
112
113impl FromIterator<(StorageKey, StorageValue)> for BasicExternalities {
114 fn from_iter<I: IntoIterator<Item=(StorageKey, StorageValue)>>(iter: I) -> Self {
115 let mut t = Self::default();
116 t.inner.top.extend(iter);
117 t
118 }
119}
120
121impl Default for BasicExternalities {
122 fn default() -> Self { Self::new(Default::default()) }
123}
124
125impl From<BTreeMap<StorageKey, StorageValue>> for BasicExternalities {
126 fn from(hashmap: BTreeMap<StorageKey, StorageValue>) -> Self {
127 BasicExternalities {
128 inner: Storage {
129 top: hashmap,
130 children_default: Default::default(),
131 },
132 extensions: Default::default(),
133 }
134 }
135}
136
137impl Externalities for BasicExternalities {
138 fn set_offchain_storage(&mut self, _key: &[u8], _value: Option<&[u8]>) {}
139
140 fn storage(&self, key: &[u8]) -> Option<StorageValue> {
141 self.inner.top.get(key).cloned()
142 }
143
144 fn storage_hash(&self, key: &[u8]) -> Option<Vec<u8>> {
145 self.storage(key).map(|v| Blake2Hasher::hash(&v).encode())
146 }
147
148 fn child_storage(
149 &self,
150 child_info: &ChildInfo,
151 key: &[u8],
152 ) -> Option<StorageValue> {
153 self.inner.children_default.get(child_info.storage_key())
154 .and_then(|child| child.data.get(key)).cloned()
155 }
156
157 fn child_storage_hash(
158 &self,
159 child_info: &ChildInfo,
160 key: &[u8],
161 ) -> Option<Vec<u8>> {
162 self.child_storage(child_info, key).map(|v| Blake2Hasher::hash(&v).encode())
163 }
164
165 fn next_storage_key(&self, key: &[u8]) -> Option<StorageKey> {
166 let range = (Bound::Excluded(key), Bound::Unbounded);
167 self.inner.top.range::<[u8], _>(range).next().map(|(k, _)| k).cloned()
168 }
169
170 fn next_child_storage_key(
171 &self,
172 child_info: &ChildInfo,
173 key: &[u8],
174 ) -> Option<StorageKey> {
175 let range = (Bound::Excluded(key), Bound::Unbounded);
176 self.inner.children_default.get(child_info.storage_key())
177 .and_then(|child| child.data.range::<[u8], _>(range).next().map(|(k, _)| k).cloned())
178 }
179
180 fn place_storage(&mut self, key: StorageKey, maybe_value: Option<StorageValue>) {
181 if is_child_storage_key(&key) {
182 warn!(target: "trie", "Refuse to set child storage key via main storage");
183 return;
184 }
185
186 match maybe_value {
187 Some(value) => { self.inner.top.insert(key, value); }
188 None => { self.inner.top.remove(&key); }
189 }
190 }
191
192 fn place_child_storage(
193 &mut self,
194 child_info: &ChildInfo,
195 key: StorageKey,
196 value: Option<StorageValue>,
197 ) {
198 let child_map = self.inner.children_default.entry(child_info.storage_key().to_vec())
199 .or_insert_with(|| StorageChild {
200 data: Default::default(),
201 child_info: child_info.to_owned(),
202 });
203 if let Some(value) = value {
204 child_map.data.insert(key, value);
205 } else {
206 child_map.data.remove(&key);
207 }
208 }
209
210 fn kill_child_storage(
211 &mut self,
212 child_info: &ChildInfo,
213 _limit: Option<u32>,
214 ) -> bool {
215 self.inner.children_default.remove(child_info.storage_key());
216 true
217 }
218
219 fn clear_prefix(&mut self, prefix: &[u8]) {
220 if is_child_storage_key(prefix) {
221 warn!(
222 target: "trie",
223 "Refuse to clear prefix that is part of child storage key via main storage"
224 );
225 return;
226 }
227
228 let to_remove = self.inner.top.range::<[u8], _>((Bound::Included(prefix), Bound::Unbounded))
229 .map(|(k, _)| k)
230 .take_while(|k| k.starts_with(prefix))
231 .cloned()
232 .collect::<Vec<_>>();
233
234 for key in to_remove {
235 self.inner.top.remove(&key);
236 }
237 }
238
239 fn clear_child_prefix(
240 &mut self,
241 child_info: &ChildInfo,
242 prefix: &[u8],
243 ) {
244 if let Some(child) = self.inner.children_default.get_mut(child_info.storage_key()) {
245 let to_remove = child.data.range::<[u8], _>((Bound::Included(prefix), Bound::Unbounded))
246 .map(|(k, _)| k)
247 .take_while(|k| k.starts_with(prefix))
248 .cloned()
249 .collect::<Vec<_>>();
250
251 for key in to_remove {
252 child.data.remove(&key);
253 }
254 }
255 }
256
257 fn storage_append(
258 &mut self,
259 key: Vec<u8>,
260 value: Vec<u8>,
261 ) {
262 let current = self.inner.top.entry(key).or_default();
263 crate::ext::StorageAppend::new(current).append(value);
264 }
265
266 fn storage_root(&mut self) -> Vec<u8> {
267 let mut top = self.inner.top.clone();
268 let prefixed_keys: Vec<_> = self.inner.children_default.iter().map(|(_k, v)| {
269 (v.child_info.prefixed_storage_key(), v.child_info.clone())
270 }).collect();
271 let empty_hash = empty_child_trie_root::<Layout<Blake2Hasher>>();
275 for (prefixed_storage_key, child_info) in prefixed_keys {
276 let child_root = self.child_storage_root(&child_info);
277 if &empty_hash[..] == &child_root[..] {
278 top.remove(prefixed_storage_key.as_slice());
279 } else {
280 top.insert(prefixed_storage_key.into_inner(), child_root);
281 }
282 }
283
284 Layout::<Blake2Hasher>::tetsy_trie_root(self.inner.top.clone()).as_ref().into()
285 }
286
287 fn child_storage_root(
288 &mut self,
289 child_info: &ChildInfo,
290 ) -> Vec<u8> {
291 if let Some(child) = self.inner.children_default.get(child_info.storage_key()) {
292 let delta = child.data.iter().map(|(k, v)| (k.as_ref(), Some(v.as_ref())));
293 crate::in_memory_backend::new_in_mem::<Blake2Hasher>()
294 .child_storage_root(&child.child_info, delta).0
295 } else {
296 empty_child_trie_root::<Layout<Blake2Hasher>>()
297 }.encode()
298 }
299
300 fn storage_changes_root(&mut self, _parent: &[u8]) -> Result<Option<Vec<u8>>, ()> {
301 Ok(None)
302 }
303
304 fn storage_start_transaction(&mut self) {
305 unimplemented!("Transactions are not supported by BasicExternalities");
306 }
307
308 fn storage_rollback_transaction(&mut self) -> Result<(), ()> {
309 unimplemented!("Transactions are not supported by BasicExternalities");
310 }
311
312 fn storage_commit_transaction(&mut self) -> Result<(), ()> {
313 unimplemented!("Transactions are not supported by BasicExternalities");
314 }
315
316 fn wipe(&mut self) {}
317
318 fn commit(&mut self) {}
319
320 fn read_write_count(&self) -> (u32, u32, u32, u32) {
321 unimplemented!("read_write_count is not supported in Basic")
322 }
323
324 fn reset_read_write_count(&mut self) {
325 unimplemented!("reset_read_write_count is not supported in Basic")
326 }
327
328 fn get_whitelist(&self) -> Vec<TrackedStorageKey> {
329 unimplemented!("get_whitelist is not supported in Basic")
330 }
331
332 fn set_whitelist(&mut self, _: Vec<TrackedStorageKey>) {
333 unimplemented!("set_whitelist is not supported in Basic")
334 }
335}
336
337impl externalities::ExtensionStore for BasicExternalities {
338 fn extension_by_type_id(&mut self, type_id: TypeId) -> Option<&mut dyn Any> {
339 self.extensions.get_mut(type_id)
340 }
341
342 fn register_extension_with_type_id(
343 &mut self,
344 type_id: TypeId,
345 extension: Box<dyn externalities::Extension>,
346 ) -> Result<(), externalities::Error> {
347 self.extensions.register_with_type_id(type_id, extension)
348 }
349
350 fn deregister_extension_by_type_id(&mut self, type_id: TypeId) -> Result<(), externalities::Error> {
351 if self.extensions.deregister(type_id) {
352 Ok(())
353 } else {
354 Err(externalities::Error::ExtensionIsNotRegistered(type_id))
355 }
356 }
357}
358
359#[cfg(test)]
360mod tests {
361 use super::*;
362 use tet_core::map;
363 use tet_core::storage::{Storage, StorageChild};
364 use tet_core::storage::well_known_keys::CODE;
365 use hex_literal::hex;
366
367 #[test]
368 fn commit_should_work() {
369 let mut ext = BasicExternalities::default();
370 ext.set_storage(b"doe".to_vec(), b"reindeer".to_vec());
371 ext.set_storage(b"dog".to_vec(), b"puppy".to_vec());
372 ext.set_storage(b"dogglesworth".to_vec(), b"cat".to_vec());
373 const ROOT: [u8; 32] = hex!("39245109cef3758c2eed2ccba8d9b370a917850af3824bc8348d505df2c298fa");
374
375 assert_eq!(&ext.storage_root()[..], &ROOT);
376 }
377
378 #[test]
379 fn set_and_retrieve_code() {
380 let mut ext = BasicExternalities::default();
381
382 let code = vec![1, 2, 3];
383 ext.set_storage(CODE.to_vec(), code.clone());
384
385 assert_eq!(&ext.storage(CODE).unwrap(), &code);
386 }
387
388 #[test]
389 fn children_works() {
390 let child_info = ChildInfo::new_default(b"storage_key");
391 let child_info = &child_info;
392 let mut ext = BasicExternalities::new(Storage {
393 top: Default::default(),
394 children_default: map![
395 child_info.storage_key().to_vec() => StorageChild {
396 data: map![ b"doe".to_vec() => b"reindeer".to_vec() ],
397 child_info: child_info.to_owned(),
398 }
399 ]
400 });
401
402 assert_eq!(ext.child_storage(child_info, b"doe"), Some(b"reindeer".to_vec()));
403
404 ext.set_child_storage(child_info, b"dog".to_vec(), b"puppy".to_vec());
405 assert_eq!(ext.child_storage(child_info, b"dog"), Some(b"puppy".to_vec()));
406
407 ext.clear_child_storage(child_info, b"dog");
408 assert_eq!(ext.child_storage(child_info, b"dog"), None);
409
410 ext.kill_child_storage(child_info, None);
411 assert_eq!(ext.child_storage(child_info, b"doe"), None);
412 }
413
414 #[test]
415 fn basic_externalities_is_empty() {
416 let storage = BasicExternalities::new_empty().into_storages();
418 assert!(storage.top.is_empty());
419 assert!(storage.children_default.is_empty());
420 }
421}