referrerpolicy=no-referrer-when-downgrade

sp_state_machine/
in_memory_backend.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: Apache-2.0
5
6// Licensed under the Apache License, Version 2.0 (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10// 	http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License.
17
18//! State machine in memory backend.
19
20use crate::{
21	backend::Backend, trie_backend::TrieBackend, StorageCollection, StorageKey, StorageValue,
22	TrieBackendBuilder,
23};
24use alloc::{collections::BTreeMap, vec::Vec};
25use codec::Codec;
26use hash_db::Hasher;
27use sp_core::storage::{ChildInfo, StateVersion, Storage};
28use sp_trie::{empty_trie_root, LayoutV1, PrefixedMemoryDB, RandomState};
29
30#[cfg(feature = "std")]
31use std::collections::HashMap as MapType;
32
33#[cfg(not(feature = "std"))]
34use alloc::collections::BTreeMap as MapType;
35
36/// Create a new empty instance of in-memory backend.
37pub fn new_in_mem<H>() -> TrieBackend<PrefixedMemoryDB<H>, H>
38where
39	H: Hasher,
40	H::Out: Codec + Ord,
41{
42	// V1 is same as V0 for an empty trie.
43	TrieBackendBuilder::new(
44		PrefixedMemoryDB::with_hasher(RandomState::default()),
45		empty_trie_root::<LayoutV1<H>>(),
46	)
47	.build()
48}
49
50impl<H: Hasher> TrieBackend<PrefixedMemoryDB<H>, H>
51where
52	H::Out: Codec + Ord,
53{
54	/// Copy the state, with applied updates
55	pub fn update<T: IntoIterator<Item = (Option<ChildInfo>, StorageCollection)>>(
56		&self,
57		changes: T,
58		state_version: StateVersion,
59	) -> Self {
60		let mut clone = self.clone();
61		clone.insert(changes, state_version);
62		clone
63	}
64
65	/// Insert values into backend trie.
66	pub fn insert<T: IntoIterator<Item = (Option<ChildInfo>, StorageCollection)>>(
67		&mut self,
68		changes: T,
69		state_version: StateVersion,
70	) {
71		let (top, child) = changes.into_iter().partition::<Vec<_>, _>(|v| v.0.is_none());
72		let (root, transaction) = self.full_storage_root(
73			top.iter().flat_map(|(_, v)| v).map(|(k, v)| (&k[..], v.as_deref())),
74			child.iter().filter_map(|v| {
75				v.0.as_ref().map(|c| (c, v.1.iter().map(|(k, v)| (&k[..], v.as_deref()))))
76			}),
77			state_version,
78		);
79
80		self.apply_transaction(root, transaction);
81	}
82
83	/// Merge trie nodes into this backend.
84	pub fn update_backend(&self, root: H::Out, changes: PrefixedMemoryDB<H>) -> Self {
85		let mut clone = self.backend_storage().clone();
86		clone.consolidate(changes);
87		TrieBackendBuilder::new(clone, root).build()
88	}
89
90	/// Apply the given transaction to this backend and set the root to the given value.
91	pub fn apply_transaction(&mut self, root: H::Out, transaction: PrefixedMemoryDB<H>) {
92		let mut storage = core::mem::take(self).into_storage();
93
94		storage.consolidate(transaction);
95		*self = TrieBackendBuilder::new(storage, root).build();
96	}
97
98	/// Compare with another in-memory backend.
99	pub fn eq(&self, other: &Self) -> bool {
100		self.root() == other.root()
101	}
102}
103
104impl<H: Hasher> Clone for TrieBackend<PrefixedMemoryDB<H>, H>
105where
106	H::Out: Codec + Ord,
107{
108	fn clone(&self) -> Self {
109		TrieBackendBuilder::new(self.backend_storage().clone(), *self.root()).build()
110	}
111}
112
113impl<H> Default for TrieBackend<PrefixedMemoryDB<H>, H>
114where
115	H: Hasher,
116	H::Out: Codec + Ord,
117{
118	fn default() -> Self {
119		new_in_mem()
120	}
121}
122
123impl<H: Hasher> From<(MapType<Option<ChildInfo>, BTreeMap<StorageKey, StorageValue>>, StateVersion)>
124	for TrieBackend<PrefixedMemoryDB<H>, H>
125where
126	H::Out: Codec + Ord,
127{
128	fn from(
129		(inner, state_version): (
130			MapType<Option<ChildInfo>, BTreeMap<StorageKey, StorageValue>>,
131			StateVersion,
132		),
133	) -> Self {
134		let mut backend = new_in_mem();
135		backend.insert(
136			inner
137				.into_iter()
138				.map(|(k, m)| (k, m.into_iter().map(|(k, v)| (k, Some(v))).collect())),
139			state_version,
140		);
141		backend
142	}
143}
144
145#[cfg(feature = "std")]
146impl<H: Hasher> From<(Storage, StateVersion)> for TrieBackend<PrefixedMemoryDB<H>, H>
147where
148	H::Out: Codec + Ord,
149{
150	fn from((inners, state_version): (Storage, StateVersion)) -> Self {
151		let mut inner: MapType<Option<ChildInfo>, BTreeMap<StorageKey, StorageValue>> = inners
152			.children_default
153			.into_values()
154			.map(|c| (Some(c.child_info), c.data))
155			.collect();
156		inner.insert(None, inners.top);
157		(inner, state_version).into()
158	}
159}
160
161impl<H: Hasher> From<(BTreeMap<StorageKey, StorageValue>, StateVersion)>
162	for TrieBackend<PrefixedMemoryDB<H>, H>
163where
164	H::Out: Codec + Ord,
165{
166	fn from((inner, state_version): (BTreeMap<StorageKey, StorageValue>, StateVersion)) -> Self {
167		let mut expanded = MapType::new();
168		expanded.insert(None, inner);
169		(expanded, state_version).into()
170	}
171}
172
173impl<H: Hasher> From<(Vec<(Option<ChildInfo>, StorageCollection)>, StateVersion)>
174	for TrieBackend<PrefixedMemoryDB<H>, H>
175where
176	H::Out: Codec + Ord,
177{
178	fn from(
179		(inner, state_version): (Vec<(Option<ChildInfo>, StorageCollection)>, StateVersion),
180	) -> Self {
181		let mut expanded: MapType<Option<ChildInfo>, BTreeMap<StorageKey, StorageValue>> =
182			MapType::new();
183		for (child_info, key_values) in inner {
184			let entry = expanded.entry(child_info).or_default();
185			for (key, value) in key_values {
186				if let Some(value) = value {
187					entry.insert(key, value);
188				}
189			}
190		}
191		(expanded, state_version).into()
192	}
193}
194
195#[cfg(test)]
196mod tests {
197	use super::*;
198	use crate::backend::{AsTrieBackend, Backend};
199	use sp_core::storage::StateVersion;
200	use sp_runtime::traits::BlakeTwo256;
201
202	/// Assert in memory backend with only child trie keys works as trie backend.
203	#[test]
204	fn in_memory_with_child_trie_only() {
205		let state_version = StateVersion::default();
206		let storage = new_in_mem::<BlakeTwo256>();
207		let child_info = ChildInfo::new_default(b"1");
208		let child_info = &child_info;
209		let storage = storage.update(
210			vec![(Some(child_info.clone()), vec![(b"2".to_vec(), Some(b"3".to_vec()))])],
211			state_version,
212		);
213		let trie_backend = storage.as_trie_backend();
214		assert_eq!(trie_backend.child_storage(child_info, b"2").unwrap(), Some(b"3".to_vec()));
215		let storage_key = child_info.prefixed_storage_key();
216		assert!(trie_backend.storage(storage_key.as_slice()).unwrap().is_some());
217	}
218
219	#[test]
220	fn insert_multiple_times_child_data_works() {
221		let state_version = StateVersion::default();
222		let mut storage = new_in_mem::<BlakeTwo256>();
223		let child_info = ChildInfo::new_default(b"1");
224
225		storage.insert(
226			vec![(Some(child_info.clone()), vec![(b"2".to_vec(), Some(b"3".to_vec()))])],
227			state_version,
228		);
229		storage.insert(
230			vec![(Some(child_info.clone()), vec![(b"1".to_vec(), Some(b"3".to_vec()))])],
231			state_version,
232		);
233
234		assert_eq!(storage.child_storage(&child_info, &b"2"[..]), Ok(Some(b"3".to_vec())));
235		assert_eq!(storage.child_storage(&child_info, &b"1"[..]), Ok(Some(b"3".to_vec())));
236	}
237}