Trie integrated benchmark (#5616)

substrate/bin/node/bench/src/generator.rs

@@ -0,0 +1,107 @@
+// Copyright 2020 Parity Technologies (UK) Ltd.
+// This file is part of Substrate.
+
+// Substrate is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Substrate is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
+
+use std::{collections::HashMap, sync::Arc};
+
+use kvdb::KeyValueDB;
+use node_primitives::Hash;
+use sp_trie::{DBValue, trie_types::TrieDBMut, TrieMut};
+use hash_db::{HashDB, AsHashDB, Prefix, Hasher as _};
+
+type Hasher = sp_core::Blake2Hasher;
+
+/// Will fill your database `db` with trie data from `key_values` and
+/// return root.
+pub fn generate_trie(
+	db: Arc<dyn KeyValueDB>,
+	key_values: impl IntoIterator<Item=(Vec<u8>, Vec<u8>)>,
+) -> Hash {
+	let mut root = Hash::default();
+
+	let (db, overlay) = {
+		let mut overlay = HashMap::new();
+		overlay.insert(
+			hex::decode("03170a2e7597b7b7e3d84c05391d139a62b157e78786d8c082f29dcf4c111314").expect("null key is valid"),
+			Some(vec![0]),
+		);
+		let mut trie_generator = TrieGenerator { db, overlay: &mut overlay };
+		{
+			let mut trie_db = TrieDBMut::new(&mut trie_generator, &mut root);
+
+			for (key, value) in key_values {
+				trie_db.insert(&key, &value).expect("trie insertion failed");
+			}
+
+			trie_db.commit();
+		}
+		( trie_generator.db, overlay )
+	};
+
+	let mut transaction = db.transaction();
+	for (key, value) in overlay.into_iter() {
+		match value {
+			Some(value) => transaction.put(0, &key[..], &value[..]),
+			None => transaction.delete(0, &key[..]),
+		}
+	}
+	db.write(transaction).expect("Failed to write transaction");
+
+	root
+}
+
+/// Immutable generated trie database with root.
+struct TrieGenerator<'a> {
+	db: Arc<dyn KeyValueDB>,
+	overlay: &'a mut HashMap<Vec<u8>, Option<Vec<u8>>>,
+}
+
+impl<'a> AsHashDB<Hasher, DBValue> for TrieGenerator<'a> {
+	fn as_hash_db(&self) -> &dyn hash_db::HashDB<Hasher, DBValue> { &*self }
+
+	fn as_hash_db_mut<'b>(&'b mut self) -> &'b mut (dyn HashDB<Hasher, DBValue> + 'b) {
+		&mut *self
+	}
+}
+
+impl<'a> HashDB<Hasher, DBValue> for TrieGenerator<'a> {
+	fn get(&self, key: &Hash, prefix: Prefix) -> Option<DBValue> {
+		let key = sp_trie::prefixed_key::<Hasher>(key, prefix);
+		if let Some(value) = self.overlay.get(&key) {
+			return value.clone();
+		}
+		self.db.get(0, &key).expect("Database backend error")
+	}
+
+	fn contains(&self, hash: &Hash, prefix: Prefix) -> bool {
+		self.get(hash, prefix).is_some()
+	}
+
+	fn insert(&mut self, prefix: Prefix, value: &[u8]) -> Hash {
+		let key = Hasher::hash(value);
+		self.emplace(key, prefix, value.to_vec());
+		key
+	}
+
+	fn emplace(&mut self, key: Hash, prefix: Prefix, value: DBValue) {
+		let key = sp_trie::prefixed_key::<Hasher>(&key, prefix);
+		self.overlay.insert(key, Some(value));
+	}
+
+	fn remove(&mut self, key: &Hash, prefix: Prefix) {
+		let key = sp_trie::prefixed_key::<Hasher>(key, prefix);
+		self.overlay.insert(key, None);
+	}
+}