1use 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
36pub fn new_in_mem<H>() -> TrieBackend<PrefixedMemoryDB<H>, H>
38where
39 H: Hasher,
40 H::Out: Codec + Ord,
41{
42 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 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 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 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 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 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 #[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}