1#[cfg(not(feature = "std"))]
2extern crate alloc;
3
4use crate::{
5 evm_chain_id_storage_key, evm_contract_creation_allowed_by_storage_key,
6 self_domain_id_storage_key, DomainId, PermissionedActionAllowedBy,
7};
8#[cfg(not(feature = "std"))]
9use alloc::vec::Vec;
10use domain_runtime_primitives::{EVMChainId, EthereumAccountId};
11use hash_db::Hasher;
12use parity_scale_codec::{Codec, Decode, Encode};
13use scale_info::TypeInfo;
14use sp_core::storage::{well_known_keys, ChildInfo};
15#[cfg(feature = "std")]
16use sp_core::storage::{Storage, StorageChild};
17use sp_runtime::StateVersion;
18use sp_state_machine::{Backend, TrieBackend, TrieBackendBuilder};
19use sp_std::collections::btree_map::BTreeMap;
20use sp_trie::{empty_trie_root, LayoutV0, MemoryDB};
21
22fn new_in_mem<H>() -> TrieBackend<MemoryDB<H>, H>
27where
28 H: Hasher,
29 H::Out: Codec + Ord,
30{
31 let db = MemoryDB::default();
32 TrieBackendBuilder::new(db, empty_trie_root::<LayoutV0<H>>()).build()
34}
35
36#[derive(Default, Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Encode, Decode, TypeInfo)]
38pub struct StorageKey(pub Vec<u8>);
39
40#[derive(Default, Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Encode, Decode, TypeInfo)]
42pub struct StorageData(pub Vec<u8>);
43
44type GenesisStorage = BTreeMap<StorageKey, StorageData>;
45
46#[derive(Default, Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Encode, Decode, TypeInfo)]
48pub struct RawGenesis {
49 top: GenesisStorage,
50 children_default: BTreeMap<StorageKey, GenesisStorage>,
51}
52
53impl RawGenesis {
54 pub fn set_domain_id(&mut self, domain_id: DomainId) {
55 let _ = self.top.insert(
56 self_domain_id_storage_key(),
57 StorageData(domain_id.encode()),
58 );
59 }
60
61 pub fn set_evm_chain_id(&mut self, chain_id: EVMChainId) {
62 let _ = self
63 .top
64 .insert(evm_chain_id_storage_key(), StorageData(chain_id.encode()));
65 }
66
67 pub fn set_evm_contract_creation_allowed_by(
68 &mut self,
69 contract_creation_allowed_by: &PermissionedActionAllowedBy<EthereumAccountId>,
70 ) {
71 let _ = self.top.insert(
72 evm_contract_creation_allowed_by_storage_key(),
73 StorageData(contract_creation_allowed_by.encode()),
74 );
75 }
76
77 pub fn set_top_storages(&mut self, storages: Vec<(StorageKey, StorageData)>) {
78 for (k, v) in storages {
79 let _ = self.top.insert(k, v);
80 }
81 }
82
83 fn set_runtime_code(&mut self, code: Vec<u8>) {
84 let _ = self.top.insert(
85 StorageKey(well_known_keys::CODE.to_vec()),
86 StorageData(code),
87 );
88 }
89
90 pub fn get_runtime_code(&self) -> Option<&[u8]> {
91 self.top
92 .get(&StorageKey(well_known_keys::CODE.to_vec()))
93 .map(|sd| sd.0.as_ref())
94 }
95
96 pub fn take_runtime_code(&mut self) -> Option<Vec<u8>> {
97 self.top
98 .remove(&StorageKey(well_known_keys::CODE.to_vec()))
99 .map(|sd| sd.0)
100 }
101
102 pub fn state_root<H>(&self, state_version: StateVersion) -> H::Out
103 where
104 H: Hasher,
105 H::Out: Codec + Ord,
106 {
107 let backend = new_in_mem::<H>();
108
109 let child_infos: Vec<_> = self
113 .children_default
114 .keys()
115 .map(|k| ChildInfo::new_default(k.0.as_slice()))
116 .collect();
117 let child_delta = child_infos.iter().zip(
118 self.children_default
119 .values()
120 .map(|v| v.iter().map(|(k, v)| (&k.0[..], Some(&v.0[..])))),
121 );
122
123 let (root, _) = backend.full_storage_root(
124 self.top.iter().map(|(k, v)| (&k.0[..], Some(&v.0[..]))),
125 child_delta,
126 state_version,
127 );
128
129 root
130 }
131
132 pub fn dummy(code: Vec<u8>) -> Self {
133 let mut raw_genesis = Self::default();
134 raw_genesis.set_runtime_code(code);
135 raw_genesis
136 }
137}
138
139#[cfg(feature = "std")]
140impl RawGenesis {
141 pub fn from_storage(storage: Storage) -> Self {
145 let top = storage
146 .top
147 .into_iter()
148 .map(|(k, v)| (StorageKey(k), StorageData(v)))
149 .collect();
150
151 let children_default = storage
152 .children_default
153 .into_iter()
154 .map(|(k, child)| {
155 (
156 StorageKey(k),
157 child
158 .data
159 .into_iter()
160 .map(|(k, v)| (StorageKey(k), StorageData(v)))
161 .collect(),
162 )
163 })
164 .collect();
165
166 RawGenesis {
167 top,
168 children_default,
169 }
170 }
171
172 pub fn into_storage(self) -> Storage {
176 let RawGenesis {
177 top: map,
178 children_default: children_map,
179 } = self;
180 let mut storage = Storage::default();
181
182 storage.top.extend(map.into_iter().map(|(k, v)| (k.0, v.0)));
183
184 children_map.into_iter().for_each(|(k, v)| {
185 let child_info = ChildInfo::new_default(k.0.as_slice());
186 storage
187 .children_default
188 .entry(k.0)
189 .or_insert_with(|| StorageChild {
190 data: Default::default(),
191 child_info,
192 })
193 .data
194 .extend(v.into_iter().map(|(k, v)| (k.0, v.0)));
195 });
196
197 storage
198 }
199}