Introduce a custom storage impl: Double Map

substrate/runtime/staking/src/account_db.rs

@@ -20,6 +20,7 @@ use rstd::prelude::*;
 use rstd::cell::RefCell;
 use rstd::collections::btree_map::{BTreeMap, Entry};
 use runtime_support::StorageMap;
+use double_map::StorageDoubleMap;
 use super::*;
 
 pub struct ChangeEntry<T: Trait> {
@@ -61,7 +62,7 @@ pub trait AccountDb<T: Trait> {
 pub struct DirectAccountDb;
 impl<T: Trait> AccountDb<T> for DirectAccountDb {
 	fn get_storage(&self, account: &T::AccountId, location: &[u8]) -> Option<Vec<u8>> {
-		<StorageOf<T>>::get(&(account.clone(), location.to_vec()))
+		<StorageOf<T>>::get(account.clone(), location.to_vec())
 	}
 	fn get_code(&self, account: &T::AccountId) -> Vec<u8> {
 		<CodeOf<T>>::get(account)
@@ -105,9 +106,9 @@ impl<T: Trait> AccountDb<T> for DirectAccountDb {
 			}
 			for (k, v) in changed.storage.into_iter() {
 				if let Some(value) = v {
-					<StorageOf<T>>::insert((address.clone(), k), &value);
+					<StorageOf<T>>::insert(address.clone(), k, value);
 				} else {
-					<StorageOf<T>>::remove((address.clone(), k));
+					<StorageOf<T>>::remove(address.clone(), k);
 				}
 			}
 		}

substrate/runtime/staking/src/double_map.rs

@@ -0,0 +1,69 @@
+// Copyright 2017 Parity Technologies (UK) Ltd.
+// This file is part of Substrate Demo.
+
+// Substrate Demo 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 Demo 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 Demo.  If not, see <http://www.gnu.org/licenses/>.
+
+//!
+
+use rstd::prelude::*;
+use codec::Slicable;
+use runtime_support::storage::unhashed;
+use runtime_io::{blake2_256, twox_128};
+
+/// Returns only a first part of the storage key.
+///
+/// Hashed by XX.
+fn first_part_of_key<M: StorageDoubleMap + ?Sized>(k1: M::Key1) -> [u8; 16] {
+	let mut raw_prefix = Vec::new();
+	raw_prefix.extend(M::PREFIX);
+	raw_prefix.extend(Slicable::encode(&k1));
+	twox_128(&raw_prefix)
+}
+
+/// Returns a compound key that consist of the two parts: (prefix, `k1`) and `k2`.
+///
+/// The first part is hased by XX and then concatenated with a blake2 hash of `k2`.
+fn full_key<M: StorageDoubleMap + ?Sized>(k1: M::Key1, k2: M::Key2) -> Vec<u8> {
+	let first_part = first_part_of_key::<M>(k1);
+	let second_part = blake2_256(&Slicable::encode(&k2));
+
+	let mut k = Vec::new();
+	k.extend(&first_part);
+	k.extend(&second_part);
+	k
+}
+
+pub trait StorageDoubleMap {
+	type Key1: Slicable;
+	type Key2: Slicable;
+	type Value: Slicable + Default;
+
+	const PREFIX: &'static [u8];
+
+	fn insert(k1: Self::Key1, k2: Self::Key2, val: Self::Value) {
+		unhashed::put(&full_key::<Self>(k1, k2)[..], &val);
+	}
+
+	fn remove(k1: Self::Key1, k2: Self::Key2) {
+		unhashed::kill(&full_key::<Self>(k1, k2)[..]);
+	}
+
+	fn get(k1: Self::Key1, k2: Self::Key2) -> Option<Self::Value> {
+		unhashed::get(&full_key::<Self>(k1, k2)[..])
+	}
+
+	fn remove_prefix(k1: Self::Key1) {
+		unhashed::kill_prefix(&first_part_of_key::<Self>(k1))
+	}
+}

substrate/runtime/staking/src/lib.rs

@@ -63,6 +63,7 @@ mod mock;
 mod tests;
 mod genesis_config;
 mod account_db;
+mod double_map;
 
 #[cfg(feature = "std")]
 pub use genesis_config::GenesisConfig;
@@ -241,7 +242,15 @@ decl_storage! {
 	// storage items with a particular prefix otherwise we'll suffer leakage with the removal
 	// of smart contracts.
 //	pub StorageOf: b"sta:sto:" => map [ T::AccountId => map(blake2) Vec<u8> => Vec<u8> ];
-	pub StorageOf: b"sta:sto:" => map [ (T::AccountId, Vec<u8>) => Vec<u8> ];
+//	pub StorageOf: b"sta:sto:" => map [ (T::AccountId, Vec<u8>) => Vec<u8> ];
+}
+
+pub(crate) struct StorageOf<T>(::rstd::marker::PhantomData<T>);
+impl<T: Trait> double_map::StorageDoubleMap for StorageOf<T> {
+	type Key1 = T::AccountId;
+	type Key2 = Vec<u8>;
+	type Value = Vec<u8>;
+	const PREFIX: &'static [u8] = b"sta:sto:";
 }
 
 enum NewAccountOutcome {
@@ -623,7 +632,7 @@ impl<T: Trait> Module<T> {
 			.map(|v| (Self::voting_balance(&v) + Self::nomination_balance(&v), v))
 			.collect::<Vec<_>>();
 		intentions.sort_unstable_by(|&(ref b1, _), &(ref b2, _)| b2.cmp(&b1));
-		
+
 		<StakeThreshold<T>>::put(
 			if intentions.len() > 0 {
 				let i = (<ValidatorCount<T>>::get() as usize).min(intentions.len() - 1);