1use crate::{Backend, OverlayedChanges, StorageKey, StorageValue};
21use alloc::{boxed::Box, collections::BTreeMap, vec::Vec};
22use codec::Encode;
23use core::{
24 any::{Any, TypeId},
25 iter::FromIterator,
26};
27use hash_db::Hasher;
28use log::warn;
29use sp_core::{
30 storage::{
31 well_known_keys::is_child_storage_key, ChildInfo, StateVersion, Storage, TrackedStorageKey,
32 },
33 traits::Externalities,
34 Blake2Hasher,
35};
36use sp_externalities::{Extension, Extensions, MultiRemovalResults};
37use sp_trie::{empty_child_trie_root, LayoutV0, LayoutV1, TrieConfiguration};
38
39#[derive(Debug)]
41pub struct BasicExternalities {
42 overlay: OverlayedChanges<Blake2Hasher>,
43 extensions: Extensions,
44}
45
46impl BasicExternalities {
47 pub fn new(inner: Storage) -> Self {
49 BasicExternalities { overlay: inner.into(), extensions: Default::default() }
50 }
51
52 pub fn new_empty() -> Self {
54 Self::new(Storage::default())
55 }
56
57 pub fn insert(&mut self, k: StorageKey, v: StorageValue) {
59 self.overlay.set_storage(k, Some(v));
60 }
61
62 #[cfg(feature = "std")]
64 pub fn into_storages(mut self) -> Storage {
65 Storage {
66 top: self
67 .overlay
68 .changes_mut()
69 .filter_map(|(k, v)| v.value().map(|v| (k.to_vec(), v.to_vec())))
70 .collect(),
71 children_default: self
72 .overlay
73 .children_mut()
74 .map(|(iter, i)| {
75 (
76 i.storage_key().to_vec(),
77 sp_core::storage::StorageChild {
78 data: iter
79 .filter_map(|(k, v)| v.value().map(|v| (k.to_vec(), v.to_vec())))
80 .collect(),
81 child_info: i.clone(),
82 },
83 )
84 })
85 .collect(),
86 }
87 }
88
89 #[cfg(feature = "std")]
93 pub fn execute_with_storage<R>(
94 storage: &mut sp_core::storage::Storage,
95 f: impl FnOnce() -> R,
96 ) -> R {
97 let mut ext = Self::new(core::mem::take(storage));
98
99 let r = ext.execute_with(f);
100
101 *storage = ext.into_storages();
102
103 r
104 }
105
106 pub fn execute_with<R>(&mut self, f: impl FnOnce() -> R) -> R {
110 sp_externalities::set_and_run_with_externalities(self, f)
111 }
112
113 pub fn extensions(&mut self) -> &mut Extensions {
115 &mut self.extensions
116 }
117
118 pub fn register_extension(&mut self, ext: impl Extension) {
120 self.extensions.register(ext);
121 }
122}
123
124#[cfg(test)]
125impl PartialEq for BasicExternalities {
126 fn eq(&self, other: &Self) -> bool {
127 self.overlay
128 .changes()
129 .map(|(k, v)| (k, v.value_ref().materialize()))
130 .collect::<BTreeMap<_, _>>() ==
131 other
132 .overlay
133 .changes()
134 .map(|(k, v)| (k, v.value_ref().materialize()))
135 .collect::<BTreeMap<_, _>>() &&
136 self.overlay
137 .children()
138 .map(|(iter, i)| {
139 (
140 i,
141 iter.map(|(k, v)| (k, v.value_ref().materialize()))
142 .collect::<BTreeMap<_, _>>(),
143 )
144 })
145 .collect::<BTreeMap<_, _>>() ==
146 other
147 .overlay
148 .children()
149 .map(|(iter, i)| {
150 (
151 i,
152 iter.map(|(k, v)| (k, v.value_ref().materialize()))
153 .collect::<BTreeMap<_, _>>(),
154 )
155 })
156 .collect::<BTreeMap<_, _>>()
157 }
158}
159
160impl FromIterator<(StorageKey, StorageValue)> for BasicExternalities {
161 fn from_iter<I: IntoIterator<Item = (StorageKey, StorageValue)>>(iter: I) -> Self {
162 let mut t = Self::default();
163 iter.into_iter().for_each(|(k, v)| t.insert(k, v));
164 t
165 }
166}
167
168impl Default for BasicExternalities {
169 fn default() -> Self {
170 Self::new(Default::default())
171 }
172}
173
174impl From<BTreeMap<StorageKey, StorageValue>> for BasicExternalities {
175 fn from(map: BTreeMap<StorageKey, StorageValue>) -> Self {
176 Self::from_iter(map.into_iter())
177 }
178}
179
180impl Externalities for BasicExternalities {
181 fn set_offchain_storage(&mut self, _key: &[u8], _value: Option<&[u8]>) {}
182
183 fn storage(&mut self, key: &[u8]) -> Option<StorageValue> {
184 self.overlay.storage(key).and_then(|v| v.map(|v| v.to_vec()))
185 }
186
187 fn storage_hash(&mut self, key: &[u8]) -> Option<Vec<u8>> {
188 self.storage(key).map(|v| Blake2Hasher::hash(&v).encode())
189 }
190
191 fn child_storage(&mut self, child_info: &ChildInfo, key: &[u8]) -> Option<StorageValue> {
192 self.overlay.child_storage(child_info, key).and_then(|v| v.map(|v| v.to_vec()))
193 }
194
195 fn child_storage_hash(&mut self, child_info: &ChildInfo, key: &[u8]) -> Option<Vec<u8>> {
196 self.child_storage(child_info, key).map(|v| Blake2Hasher::hash(&v).encode())
197 }
198
199 fn next_storage_key(&mut self, key: &[u8]) -> Option<StorageKey> {
200 self.overlay.iter_after(key).find_map(|(k, v)| v.value().map(|_| k.to_vec()))
201 }
202
203 fn next_child_storage_key(&mut self, child_info: &ChildInfo, key: &[u8]) -> Option<StorageKey> {
204 self.overlay
205 .child_iter_after(child_info.storage_key(), key)
206 .find_map(|(k, v)| v.value().map(|_| k.to_vec()))
207 }
208
209 fn place_storage(&mut self, key: StorageKey, maybe_value: Option<StorageValue>) {
210 if is_child_storage_key(&key) {
211 warn!(target: "trie", "Refuse to set child storage key via main storage");
212 return
213 }
214
215 self.overlay.set_storage(key, maybe_value)
216 }
217
218 fn place_child_storage(
219 &mut self,
220 child_info: &ChildInfo,
221 key: StorageKey,
222 value: Option<StorageValue>,
223 ) {
224 self.overlay.set_child_storage(child_info, key, value);
225 }
226
227 fn kill_child_storage(
228 &mut self,
229 child_info: &ChildInfo,
230 _maybe_limit: Option<u32>,
231 _maybe_cursor: Option<&[u8]>,
232 ) -> MultiRemovalResults {
233 let count = self.overlay.clear_child_storage(child_info);
234 MultiRemovalResults { maybe_cursor: None, backend: count, unique: count, loops: count }
235 }
236
237 fn clear_prefix(
238 &mut self,
239 prefix: &[u8],
240 _maybe_limit: Option<u32>,
241 _maybe_cursor: Option<&[u8]>,
242 ) -> MultiRemovalResults {
243 if is_child_storage_key(prefix) {
244 warn!(
245 target: "trie",
246 "Refuse to clear prefix that is part of child storage key via main storage"
247 );
248 let maybe_cursor = Some(prefix.to_vec());
249 return MultiRemovalResults { maybe_cursor, backend: 0, unique: 0, loops: 0 }
250 }
251
252 let count = self.overlay.clear_prefix(prefix);
253 MultiRemovalResults { maybe_cursor: None, backend: count, unique: count, loops: count }
254 }
255
256 fn clear_child_prefix(
257 &mut self,
258 child_info: &ChildInfo,
259 prefix: &[u8],
260 _maybe_limit: Option<u32>,
261 _maybe_cursor: Option<&[u8]>,
262 ) -> MultiRemovalResults {
263 let count = self.overlay.clear_child_prefix(child_info, prefix);
264 MultiRemovalResults { maybe_cursor: None, backend: count, unique: count, loops: count }
265 }
266
267 fn storage_append(&mut self, key: Vec<u8>, element: Vec<u8>) {
268 self.overlay.append_storage(key, element, Default::default);
269 }
270
271 fn storage_root(&mut self, state_version: StateVersion) -> Vec<u8> {
272 let mut top = self
273 .overlay
274 .changes_mut()
275 .filter_map(|(k, v)| v.value().map(|v| (k.clone(), v.clone())))
276 .collect::<BTreeMap<_, _>>();
277 let empty_hash = empty_child_trie_root::<LayoutV1<Blake2Hasher>>();
281 for child_info in self.overlay.children().map(|d| d.1.clone()).collect::<Vec<_>>() {
282 let child_root = self.child_storage_root(&child_info, state_version);
283 if empty_hash[..] == child_root[..] {
284 top.remove(child_info.prefixed_storage_key().as_slice());
285 } else {
286 top.insert(child_info.prefixed_storage_key().into_inner(), child_root);
287 }
288 }
289
290 match state_version {
291 StateVersion::V0 => LayoutV0::<Blake2Hasher>::trie_root(top).as_ref().into(),
292 StateVersion::V1 => LayoutV1::<Blake2Hasher>::trie_root(top).as_ref().into(),
293 }
294 }
295
296 fn child_storage_root(
297 &mut self,
298 child_info: &ChildInfo,
299 state_version: StateVersion,
300 ) -> Vec<u8> {
301 if let Some((data, child_info)) = self.overlay.child_changes_mut(child_info.storage_key()) {
302 let delta =
303 data.into_iter().map(|(k, v)| (k.as_ref(), v.value().map(|v| v.as_slice())));
304 crate::in_memory_backend::new_in_mem::<Blake2Hasher>()
305 .child_storage_root(&child_info, delta, state_version)
306 .0
307 } else {
308 empty_child_trie_root::<LayoutV1<Blake2Hasher>>()
309 }
310 .encode()
311 }
312
313 fn storage_start_transaction(&mut self) {
314 self.overlay.start_transaction()
315 }
316
317 fn storage_rollback_transaction(&mut self) -> Result<(), ()> {
318 self.overlay.rollback_transaction().map_err(drop)
319 }
320
321 fn storage_commit_transaction(&mut self) -> Result<(), ()> {
322 self.overlay.commit_transaction().map_err(drop)
323 }
324
325 fn wipe(&mut self) {}
326
327 fn commit(&mut self) {}
328
329 fn read_write_count(&self) -> (u32, u32, u32, u32) {
330 unimplemented!("read_write_count is not supported in Basic")
331 }
332
333 fn reset_read_write_count(&mut self) {
334 unimplemented!("reset_read_write_count is not supported in Basic")
335 }
336
337 fn get_whitelist(&self) -> Vec<TrackedStorageKey> {
338 unimplemented!("get_whitelist is not supported in Basic")
339 }
340
341 fn set_whitelist(&mut self, _: Vec<TrackedStorageKey>) {
342 unimplemented!("set_whitelist is not supported in Basic")
343 }
344
345 fn get_read_and_written_keys(&self) -> Vec<(Vec<u8>, u32, u32, bool)> {
346 unimplemented!("get_read_and_written_keys is not supported in Basic")
347 }
348}
349
350impl sp_externalities::ExtensionStore for BasicExternalities {
351 fn extension_by_type_id(&mut self, type_id: TypeId) -> Option<&mut dyn Any> {
352 self.extensions.get_mut(type_id)
353 }
354
355 fn register_extension_with_type_id(
356 &mut self,
357 type_id: TypeId,
358 extension: Box<dyn sp_externalities::Extension>,
359 ) -> Result<(), sp_externalities::Error> {
360 self.extensions.register_with_type_id(type_id, extension)
361 }
362
363 fn deregister_extension_by_type_id(
364 &mut self,
365 type_id: TypeId,
366 ) -> Result<(), sp_externalities::Error> {
367 if self.extensions.deregister(type_id) {
368 Ok(())
369 } else {
370 Err(sp_externalities::Error::ExtensionIsNotRegistered(type_id))
371 }
372 }
373}
374
375#[cfg(test)]
376mod tests {
377 use super::*;
378 use sp_core::{
379 map,
380 storage::{well_known_keys::CODE, Storage, StorageChild},
381 };
382
383 #[test]
384 fn commit_should_work() {
385 let mut ext = BasicExternalities::default();
386 ext.set_storage(b"doe".to_vec(), b"reindeer".to_vec());
387 ext.set_storage(b"dog".to_vec(), b"puppy".to_vec());
388 ext.set_storage(b"dogglesworth".to_vec(), b"cat".to_vec());
389 let root = array_bytes::hex2bytes_unchecked(
390 "39245109cef3758c2eed2ccba8d9b370a917850af3824bc8348d505df2c298fa",
391 );
392
393 assert_eq!(&ext.storage_root(StateVersion::default())[..], &root);
394 }
395
396 #[test]
397 fn set_and_retrieve_code() {
398 let mut ext = BasicExternalities::default();
399
400 let code = vec![1, 2, 3];
401 ext.set_storage(CODE.to_vec(), code.clone());
402
403 assert_eq!(&ext.storage(CODE).unwrap(), &code);
404 }
405
406 #[test]
407 fn children_works() {
408 let child_info = ChildInfo::new_default(b"storage_key");
409 let child_info = &child_info;
410 let mut ext = BasicExternalities::new(Storage {
411 top: Default::default(),
412 children_default: map![
413 child_info.storage_key().to_vec() => StorageChild {
414 data: map![ b"doe".to_vec() => b"reindeer".to_vec() ],
415 child_info: child_info.to_owned(),
416 }
417 ],
418 });
419
420 assert_eq!(ext.child_storage(child_info, b"doe"), Some(b"reindeer".to_vec()));
421
422 ext.set_child_storage(child_info, b"dog".to_vec(), b"puppy".to_vec());
423 assert_eq!(ext.child_storage(child_info, b"dog"), Some(b"puppy".to_vec()));
424
425 ext.clear_child_storage(child_info, b"dog");
426 assert_eq!(ext.child_storage(child_info, b"dog"), None);
427
428 let _ = ext.kill_child_storage(child_info, None, None);
429 assert_eq!(ext.child_storage(child_info, b"doe"), None);
430 }
431
432 #[test]
433 fn kill_child_storage_returns_num_elements_removed() {
434 let child_info = ChildInfo::new_default(b"storage_key");
435 let child_info = &child_info;
436 let mut ext = BasicExternalities::new(Storage {
437 top: Default::default(),
438 children_default: map![
439 child_info.storage_key().to_vec() => StorageChild {
440 data: map![
441 b"doe".to_vec() => b"reindeer".to_vec(),
442 b"dog".to_vec() => b"puppy".to_vec(),
443 b"hello".to_vec() => b"world".to_vec(),
444 ],
445 child_info: child_info.to_owned(),
446 }
447 ],
448 });
449
450 let res = ext.kill_child_storage(child_info, None, None);
451 assert_eq!(res.deconstruct(), (None, 3, 3, 3));
452 }
453
454 #[test]
455 fn basic_externalities_is_empty() {
456 let storage = BasicExternalities::new_empty().into_storages();
458 assert!(storage.top.is_empty());
459 assert!(storage.children_default.is_empty());
460 }
461}