Contract restoration (#2502)

* partial impl

* more checks

* improve TestExternalities + tests

* make tombstone raw to derive Eq

* remove before merge

* add test

* fmt

* update test

* doc

* bump version

* fix testing

* update runtime

* Fix TrieIdGenerator

* fix dummy trie id generator

* update test

* bump version

* format

* Update core/state-machine/src/testing.rs

Co-Authored-By: cheme <emericchevalier.pro@gmail.com>

* document test

* Apply suggestions from code review

Co-Authored-By: DemiMarie-temp <50585338+DemiMarie-temp@users.noreply.github.com>

* refactor

* fix

* fmt

* address review

* impl last_write

* Fix storage size, test, remove size in tombstone

* fix

* Update srml/contract/src/lib.rs

Co-Authored-By: Sergei Pepyakin <s.pepyakin@gmail.com>

* comment

* child_storage_root as &[u8]

substrate/Cargo.lock

@@ -3318,6 +3318,7 @@ name = "srml-contract"
 version = "2.0.0"
 dependencies = [
  "assert_matches 1.3.0 (registry+https://github.com/rust-lang/crates.io-index)",
+ "hex 0.3.2 (registry+https://github.com/rust-lang/crates.io-index)",
  "hex-literal 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "parity-codec 3.5.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "parity-wasm 0.31.3 (registry+https://github.com/rust-lang/crates.io-index)",

substrate/core/state-machine/src/ext.rs

@@ -302,7 +302,7 @@ where
 				.flat_map(|map| map.1.iter().map(|(k, v)| (k.clone(), v.clone())))
 				.chain(self.overlay.prospective.children.get(storage_key)
 						.into_iter()
-						.flat_map(|map| map.1.iter().map(|(k, v)| (k.clone(), v.clone()))));
+						.flat_map(|map| map.1.clone().into_iter()));
 
 			let root = self.backend.child_storage_root(storage_key, delta).0;
 

substrate/core/state-machine/src/testing.rs

@@ -16,22 +16,27 @@
 
 //! Test implementation for Externalities.
 
-use std::collections::HashMap;
+use std::collections::{HashMap, BTreeMap};
 use std::iter::FromIterator;
+use std::marker::PhantomData;
 use hash_db::Hasher;
-use trie::trie_root;
-use crate::backend::InMemory;
-use crate::changes_trie::{compute_changes_trie_root, InMemoryStorage as ChangesTrieInMemoryStorage, AnchorBlockId};
+use crate::backend::{InMemory, Backend};
+use primitives::storage::well_known_keys::is_child_storage_key;
+use crate::changes_trie::{
+	compute_changes_trie_root, InMemoryStorage as ChangesTrieInMemoryStorage, AnchorBlockId
+};
 use primitives::storage::well_known_keys::{CHANGES_TRIE_CONFIG, CODE, HEAP_PAGES};
 use parity_codec::Encode;
 use super::{ChildStorageKey, Externalities, OverlayedChanges};
 
+const EXT_NOT_ALLOWED_TO_FAIL: &str = "Externalities not allowed to fail within runtime";
+
 /// Simple HashMap-based Externalities impl.
 pub struct TestExternalities<H: Hasher> {
-	inner: HashMap<Vec<u8>, Vec<u8>>,
+	overlay: OverlayedChanges,
+	backend: InMemory<H>,
 	changes_trie_storage: ChangesTrieInMemoryStorage<H>,
-	changes: OverlayedChanges,
-	code: Option<Vec<u8>>,
+	_hasher: PhantomData<H>,
 }
 
 impl<H: Hasher> TestExternalities<H> {
@@ -43,6 +48,7 @@ impl<H: Hasher> TestExternalities<H> {
 	/// Create a new instance of `TestExternalities`
 	pub fn new_with_code(code: &[u8], mut inner: HashMap<Vec<u8>, Vec<u8>>) -> Self {
 		let mut overlay = OverlayedChanges::default();
+
 		super::set_changes_trie_config(
 			&mut overlay,
 			inner.get(&CHANGES_TRIE_CONFIG.to_vec()).cloned(),
@@ -50,37 +56,51 @@ impl<H: Hasher> TestExternalities<H> {
 		).expect("changes trie configuration is correct in test env; qed");
 
 		inner.insert(HEAP_PAGES.to_vec(), 8u64.encode());
+		inner.insert(CODE.to_vec(), code.to_vec());
 
 		TestExternalities {
-			inner,
+			overlay,
 			changes_trie_storage: ChangesTrieInMemoryStorage::new(),
-			changes: overlay,
-			code: Some(code.to_vec()),
+			backend: inner.into(),
+			_hasher: Default::default(),
 		}
 	}
 
-	/// Insert key/value
-	pub fn insert(&mut self, k: Vec<u8>, v: Vec<u8>) -> Option<Vec<u8>> {
-		self.inner.insert(k, v)
+	/// Insert key/value into backend
+	pub fn insert(&mut self, k: Vec<u8>, v: Vec<u8>) {
+		self.backend = self.backend.update(vec![(None, k, Some(v))]);
+	}
+
+	/// Iter to all pairs in key order
+	pub fn iter_pairs_in_order(&self) -> impl Iterator<Item=(Vec<u8>, Vec<u8>)> {
+		self.backend.pairs().iter()
+			.map(|&(ref k, ref v)| (k.to_vec(), Some(v.to_vec())))
+			.chain(self.overlay.committed.top.clone().into_iter().map(|(k, v)| (k, v.value)))
+			.chain(self.overlay.prospective.top.clone().into_iter().map(|(k, v)| (k, v.value)))
+			.collect::<BTreeMap<_, _>>()
+			.into_iter()
+			.filter_map(|(k, maybe_val)| maybe_val.map(|val| (k, val)))
 	}
 }
 
 impl<H: Hasher> ::std::fmt::Debug for TestExternalities<H> {
 	fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
-		write!(f, "{:?}", self.inner)
+		write!(f, "overlay: {:?}\nbackend: {:?}", self.overlay, self.backend.pairs())
 	}
 }
 
 impl<H: Hasher> PartialEq for TestExternalities<H> {
+	/// This doesn't test if they are in the same state, only if they contains the
+	/// same data at this state
 	fn eq(&self, other: &TestExternalities<H>) -> bool {
-		self.inner.eq(&other.inner)
+		self.iter_pairs_in_order().eq(other.iter_pairs_in_order())
 	}
 }
 
 impl<H: Hasher> FromIterator<(Vec<u8>, Vec<u8>)> for TestExternalities<H> {
 	fn from_iter<I: IntoIterator<Item=(Vec<u8>, Vec<u8>)>>(iter: I) -> Self {
 		let mut t = Self::new(Default::default());
-		t.inner.extend(iter);
+		t.backend = t.backend.update(iter.into_iter().map(|(k, v)| (None, k, Some(v))).collect());
 		t
 	}
 }
@@ -91,81 +111,109 @@ impl<H: Hasher> Default for TestExternalities<H> {
 
 impl<H: Hasher> From<TestExternalities<H>> for HashMap<Vec<u8>, Vec<u8>> {
 	fn from(tex: TestExternalities<H>) -> Self {
-		tex.inner.into()
+		tex.iter_pairs_in_order().collect()
 	}
 }
 
 impl<H: Hasher> From< HashMap<Vec<u8>, Vec<u8>> > for TestExternalities<H> {
 	fn from(hashmap: HashMap<Vec<u8>, Vec<u8>>) -> Self {
-		TestExternalities {
-			inner: hashmap,
-			changes_trie_storage: ChangesTrieInMemoryStorage::new(),
-			changes: Default::default(),
-			code: None,
-		}
+		Self::from_iter(hashmap)
 	}
 }
 
-// TODO child test primitives are currently limited to `changes` (for non child the way
-// things are defined seems utterly odd to (put changes in changes but never make them
-// available for read through inner)
 impl<H: Hasher> Externalities<H> for TestExternalities<H> where H::Out: Ord {
 	fn storage(&self, key: &[u8]) -> Option<Vec<u8>> {
-		match key {
-			CODE => self.code.clone(),
-			_ => self.inner.get(key).cloned(),
-		}
+		self.overlay.storage(key).map(|x| x.map(|x| x.to_vec())).unwrap_or_else(||
+			self.backend.storage(key).expect(EXT_NOT_ALLOWED_TO_FAIL))
 	}
 
 	fn original_storage(&self, key: &[u8]) -> Option<Vec<u8>> {
-		self.storage(key)
+		self.backend.storage(key).expect(EXT_NOT_ALLOWED_TO_FAIL)
 	}
 
 	fn child_storage(&self, storage_key: ChildStorageKey<H>, key: &[u8]) -> Option<Vec<u8>> {
-		self.changes.child_storage(storage_key.as_ref(), key)?.map(Vec::from)
+		self.overlay
+			.child_storage(storage_key.as_ref(), key)
+			.map(|x| x.map(|x| x.to_vec()))
+			.unwrap_or_else(|| self.backend
+				.child_storage(storage_key.as_ref(), key)
+				.expect(EXT_NOT_ALLOWED_TO_FAIL)
+			)
 	}
 
 	fn place_storage(&mut self, key: Vec<u8>, maybe_value: Option<Vec<u8>>) {
-		self.changes.set_storage(key.clone(), maybe_value.clone());
-		match key.as_ref() {
-			CODE => self.code = maybe_value,
-			_ => {
-				match maybe_value {
-					Some(value) => { self.inner.insert(key, value); }
-					None => { self.inner.remove(&key); }
-				}
-			}
+		if is_child_storage_key(&key) {
+			panic!("Refuse to directly set child storage key");
 		}
+
+		self.overlay.set_storage(key, maybe_value);
 	}
 
-	fn place_child_storage(&mut self, storage_key: ChildStorageKey<H>, key: Vec<u8>, value: Option<Vec<u8>>) {
-		self.changes.set_child_storage(storage_key.into_owned(), key, value);
+	fn place_child_storage(
+		&mut self,
+		storage_key: ChildStorageKey<H>,
+		key: Vec<u8>,
+		value: Option<Vec<u8>>
+	) {
+		self.overlay.set_child_storage(storage_key.into_owned(), key, value);
 	}
 
 	fn kill_child_storage(&mut self, storage_key: ChildStorageKey<H>) {
-		self.changes.clear_child_storage(storage_key.as_ref());
+		let backend = &self.backend;
+		let overlay = &mut self.overlay;
+
+		overlay.clear_child_storage(storage_key.as_ref());
+		backend.for_keys_in_child_storage(storage_key.as_ref(), |key| {
+			overlay.set_child_storage(storage_key.as_ref().to_vec(), key.to_vec(), None);
+		});
 	}
 
 	fn clear_prefix(&mut self, prefix: &[u8]) {
-		self.changes.clear_prefix(prefix);
-		self.inner.retain(|key, _| !key.starts_with(prefix));
+		if is_child_storage_key(prefix) {
+			panic!("Refuse to directly clear prefix that is part of child storage key");
+		}
+
+		self.overlay.clear_prefix(prefix);
+
+		let backend = &self.backend;
+		let overlay = &mut self.overlay;
+		backend.for_keys_with_prefix(prefix, |key| {
+			overlay.set_storage(key.to_vec(), None);
+		});
 	}
 
 	fn chain_id(&self) -> u64 { 42 }
 
 	fn storage_root(&mut self) -> H::Out {
-		trie_root::<H, _, _, _>(self.inner.clone())
+		// compute and memoize
+		let delta = self.overlay.committed.top.iter().map(|(k, v)| (k.clone(), v.value.clone()))
+			.chain(self.overlay.prospective.top.iter().map(|(k, v)| (k.clone(), v.value.clone())));
+
+		self.backend.storage_root(delta).0
 	}
 
-	fn child_storage_root(&mut self, _storage_key: ChildStorageKey<H>) -> Vec<u8> {
-		vec![]
+	fn child_storage_root(&mut self, storage_key: ChildStorageKey<H>) -> Vec<u8> {
+		let storage_key = storage_key.as_ref();
+
+		let (root, _, _) = {
+			let delta = self.overlay.committed.children.get(storage_key)
+				.into_iter()
+				.flat_map(|map| map.1.iter().map(|(k, v)| (k.clone(), v.clone())))
+				.chain(self.overlay.prospective.children.get(storage_key)
+						.into_iter()
+						.flat_map(|map| map.1.clone().into_iter()));
+
+			self.backend.child_storage_root(storage_key, delta)
+		};
+		self.overlay.set_storage(storage_key.into(), Some(root.clone()));
+		root
 	}
 
 	fn storage_changes_root(&mut self, parent: H::Out, parent_num: u64) -> Option<H::Out> {
-		compute_changes_trie_root::<_, _, H>(
-			&InMemory::default(),
+		compute_changes_trie_root::<_, ChangesTrieInMemoryStorage<H>, H>(
+			&self.backend,
 			Some(&self.changes_trie_storage),
-			&self.changes,
+			&self.overlay,
 			&AnchorBlockId { hash: parent, number: parent_num },
 		).map(|(root, _)| root.clone())
 	}
@@ -187,7 +235,7 @@ mod tests {
 		ext.set_storage(b"doe".to_vec(), b"reindeer".to_vec());
 		ext.set_storage(b"dog".to_vec(), b"puppy".to_vec());
 		ext.set_storage(b"dogglesworth".to_vec(), b"cat".to_vec());
-		const ROOT: [u8; 32] = hex!("0b33ed94e74e0f8e92a55923bece1ed02d16cf424e124613ddebc53ac3eeeabe");
+		const ROOT: [u8; 32] = hex!("cc65c26c37ebd4abcdeb3f1ecd727527051620779a2f6c809bac0f8a87dbb816");
 		assert_eq!(ext.storage_root(), H256::from(ROOT));
 	}
 

substrate/node/runtime/src/lib.rs

@@ -59,7 +59,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
 	impl_name: create_runtime_str!("substrate-node"),
 	authoring_version: 10,
 	spec_version: 81,
-	impl_version: 81,
+	impl_version: 82,
 	apis: RUNTIME_API_VERSIONS,
 };
 

substrate/srml/contract/Cargo.toml

@@ -25,6 +25,7 @@ assert_matches = "1.1"
 hex-literal = "0.2.0"
 consensus = { package = "srml-consensus", path = "../consensus" }
 balances = { package = "srml-balances", path = "../balances" }
+hex = "0.3"
 
 [features]
 default = ["std"]

substrate/srml/contract/src/account_db.rs

@@ -124,6 +124,7 @@ impl<T: Trait> AccountDb<T> for DirectAccountDb {
 						trie_id: <T as Trait>::TrieIdGenerator::trie_id(&address),
 						deduct_block: <system::Module<T>>::block_number(),
 						rent_allowance: <BalanceOf<T>>::max_value(),
+						last_write: None,
 					}
 				} else {
 					// No contract exist and no code_hash provided
@@ -138,6 +139,10 @@ impl<T: Trait> AccountDb<T> for DirectAccountDb {
 					new_info.code_hash = code_hash;
 				}
 
+				if !changed.storage.is_empty() {
+					new_info.last_write = Some(<system::Module<T>>::block_number());
+				}
+
 				for (k, v) in changed.storage.into_iter() {
 					if let Some(value) = child::get_raw(&new_info.trie_id[..], &blake2_256(&k)) {
 						new_info.storage_size -= value.len() as u32;

substrate/srml/contract/src/lib.rs

@@ -96,9 +96,14 @@ use serde::{Serialize, Deserialize};
 use substrate_primitives::crypto::UncheckedFrom;
 use rstd::{prelude::*, marker::PhantomData, convert::TryFrom};
 use parity_codec::{Codec, Encode, Decode};
-use runtime_primitives::traits::{Hash, SimpleArithmetic, Bounded, StaticLookup, Zero};
+use runtime_io::blake2_256;
+use runtime_primitives::traits::{
+	Hash, SimpleArithmetic, Bounded, StaticLookup, Zero, MaybeSerializeDebug, Member
+};
 use srml_support::dispatch::{Result, Dispatchable};
-use srml_support::{Parameter, StorageMap, StorageValue, decl_module, decl_event, decl_storage, storage::child};
+use srml_support::{
+	Parameter, StorageMap, StorageValue, decl_module, decl_event, decl_storage, storage::child
+};
 use srml_support::traits::{OnFreeBalanceZero, OnUnbalanced, Currency};
 use system::{ensure_signed, RawOrigin};
 use substrate_primitives::storage::well_known_keys::CHILD_STORAGE_KEY_PREFIX;
@@ -120,6 +125,7 @@ pub trait ComputeDispatchFee<Call, Balance> {
 /// Information for managing an acocunt and its sub trie abstraction.
 /// This is the required info to cache for an account
 #[derive(Encode, Decode)]
+#[cfg_attr(feature = "std", derive(Debug))]
 pub enum ContractInfo<T: Trait> {
 	Alive(AliveContractInfo<T>),
 	Tombstone(TombstoneContractInfo<T>),
@@ -177,12 +183,14 @@ impl<T: Trait> ContractInfo<T> {
 	}
 }
 
-pub type AliveContractInfo<T> = RawAliveContractInfo<CodeHash<T>, BalanceOf<T>, <T as system::Trait>::BlockNumber>;
+pub type AliveContractInfo<T> =
+	RawAliveContractInfo<CodeHash<T>, BalanceOf<T>, <T as system::Trait>::BlockNumber>;
 
 /// Information for managing an account and its sub trie abstraction.
 /// This is the required info to cache for an account.
 // Workaround for https://github.com/rust-lang/rust/issues/26925 . Remove when sorted.
 #[derive(Encode, Decode, Clone, PartialEq, Eq)]
+#[cfg_attr(feature = "std", derive(Debug))]
 pub struct RawAliveContractInfo<CodeHash, Balance, BlockNumber> {
 	/// Unique ID for the subtree encoded as a bytes vector.
 	pub trie_id: TrieId,
@@ -190,20 +198,32 @@ pub struct RawAliveContractInfo<CodeHash, Balance, BlockNumber> {
 	pub storage_size: u32,
 	/// The code associated with a given account.
 	pub code_hash: CodeHash,
+	/// Pay rent at most up to this value.
 	pub rent_allowance: Balance,
+	/// Last block rent has been payed.
 	pub deduct_block: BlockNumber,
+	/// Last block child storage has been written.
+	pub last_write: Option<BlockNumber>,
 }
 
-#[derive(Encode, Decode)]
-pub struct TombstoneContractInfo<T: Trait>(T::Hash);
+pub type TombstoneContractInfo<T> =
+	RawTombstoneContractInfo<<T as system::Trait>::Hash, <T as system::Trait>::Hashing>;
 
-impl<T: Trait> TombstoneContractInfo<T> {
-	fn new(storage_root: Vec<u8>, storage_size: u32, code_hash: CodeHash<T>) -> Self {
+// Workaround for https://github.com/rust-lang/rust/issues/26925 . Remove when sorted.
+#[derive(Encode, Decode, PartialEq, Eq)]
+#[cfg_attr(feature = "std", derive(Debug))]
+pub struct RawTombstoneContractInfo<H, Hasher>(H, PhantomData<Hasher>);
+
+impl<H, Hasher> RawTombstoneContractInfo<H, Hasher>
+where
+	H: Member + MaybeSerializeDebug + AsRef<[u8]> + AsMut<[u8]> + Copy + Default + rstd::hash::Hash,
+	Hasher: Hash<Output=H>,
+{
+	fn new(storage_root: &[u8], code_hash: H) -> Self {
 		let mut buf = Vec::new();
 		storage_root.using_encoded(|encoded| buf.extend_from_slice(encoded));
-		storage_size.using_encoded(|encoded| buf.extend_from_slice(encoded));
 		buf.extend_from_slice(code_hash.as_ref());
-		TombstoneContractInfo(T::Hashing::hash(&buf[..]))
+		RawTombstoneContractInfo(Hasher::hash(&buf[..]), PhantomData)
 	}
 }
 
@@ -235,7 +255,10 @@ where
 	fn trie_id(account_id: &T::AccountId) -> TrieId {
 		// Note that skipping a value due to error is not an issue here.
 		// We only need uniqueness, not sequence.
-		let new_seed = <AccountCounter<T>>::mutate(|v| v.wrapping_add(1));
+		let new_seed = <AccountCounter<T>>::mutate(|v| {
+			*v = v.wrapping_add(1);
+			*v
+		});
 
 		let mut buf = Vec::new();
 		buf.extend_from_slice(account_id.as_ref());
@@ -508,6 +531,84 @@ decl_module! {
 			}
 		}
 
+		/// Allows a contract to restore a tombstone by giving its storage.
+		///
+		/// The contract that wants to restore (i.e. origin of the call, or `msg.sender` in Solidity terms) will compute a
+		/// tombstone with its storage and the given code_hash. If the computed tombstone
+		/// match the destination one, the destination contract is restored with the rent_allowance` specified,
+		/// while the origin sends all its funds to the destination and is removed.
+		fn restore_to(
+			origin,
+			dest: T::AccountId,
+			code_hash: CodeHash<T>,
+			rent_allowance: BalanceOf<T>,
+			delta: Vec<exec::StorageKey>
+		) {
+			let origin = ensure_signed(origin)?;
+
+			let mut origin_contract = <ContractInfoOf<T>>::get(&origin)
+				.and_then(|c| c.get_alive())
+				.ok_or("Cannot restore from inexisting or tombstone contract")?;
+
+			let current_block = <system::Module<T>>::block_number();
+
+			if origin_contract.last_write == Some(current_block) {
+				return Err("Origin TrieId written in the current block");
+			}
+
+			let dest_tombstone = <ContractInfoOf<T>>::get(&dest)
+				.and_then(|c| c.get_tombstone())
+				.ok_or("Cannot restore to inexisting or alive contract")?;
+
+			let last_write = if !delta.is_empty() {
+				Some(current_block)
+			} else {
+				origin_contract.last_write
+			};
+
+			let key_values_taken = delta.iter()
+				.filter_map(|key| {
+					child::get_raw(&origin_contract.trie_id, &blake2_256(key)).map(|value| {
+						child::kill(&origin_contract.trie_id, &blake2_256(key));
+						(key, value)
+					})
+				})
+				.collect::<Vec<_>>();
+
+			let tombstone = <TombstoneContractInfo<T>>::new(
+				// This operation is cheap enough because last_write (delta not included)
+				// is not this block as it has been checked earlier.
+				&runtime_io::child_storage_root(&origin_contract.trie_id)[..],
+				code_hash,
+			);
+
+			if tombstone != dest_tombstone {
+				for (key, value) in key_values_taken {
+					child::put_raw(&origin_contract.trie_id, &blake2_256(key), &value);
+				}
+
+				return Err("Tombstones don't match");
+			}
+
+			origin_contract.storage_size -= key_values_taken.iter()
+				.map(|(_, value)| value.len() as u32)
+				.sum::<u32>();
+
+			<ContractInfoOf<T>>::remove(&origin);
+			<ContractInfoOf<T>>::insert(&dest, ContractInfo::Alive(RawAliveContractInfo {
+				trie_id: origin_contract.trie_id,
+				storage_size: origin_contract.storage_size,
+				code_hash,
+				rent_allowance,
+				deduct_block: current_block,
+				last_write,
+			}));
+
+			let origin_free_balance = T::Currency::free_balance(&origin);
+			T::Currency::make_free_balance_be(&origin, <BalanceOf<T>>::zero());
+			T::Currency::deposit_creating(&dest, origin_free_balance);
+		}
+
 		fn on_finalize() {
 			<GasSpent<T>>::kill();
 		}

substrate/srml/contract/src/rent.rs

@@ -163,9 +163,8 @@ fn try_evict_or_and_pay_rent<T: Trait>(
 			// Note: this operation is heavy.
 			let child_storage_root = runtime_io::child_storage_root(&contract.trie_id);
 
-			let tombstone = TombstoneContractInfo::new(
-				child_storage_root,
-				contract.storage_size,
+			let tombstone = <TombstoneContractInfo<T>>::new(
+				&child_storage_root[..],
 				contract.code_hash,
 			);
 			<ContractInfoOf<T>>::insert(account, ContractInfo::Tombstone(tombstone));

substrate/srml/contract/src/tests.rs

@@ -34,7 +34,7 @@ use runtime_primitives::traits::{BlakeTwo256, IdentityLookup};
 use runtime_primitives::BuildStorage;
 use srml_support::{
 	assert_ok, impl_outer_dispatch, impl_outer_event, impl_outer_origin, storage::child,
-	traits::Currency, StorageMap,
+	traits::Currency, StorageMap, StorageValue
 };
 use std::sync::atomic::{AtomicUsize, Ordering};
 use substrate_primitives::storage::well_known_keys;
@@ -119,18 +119,21 @@ impl ContractAddressFor<H256, u64> for DummyContractAddressFor {
 	}
 }
 
-static KEY_COUNTER: AtomicUsize = AtomicUsize::new(0);
-
 pub struct DummyTrieIdGenerator;
 impl TrieIdGenerator<u64> for DummyTrieIdGenerator {
 	fn trie_id(account_id: &u64) -> TrieId {
 		use substrate_primitives::storage::well_known_keys;
 
+		let new_seed = <super::AccountCounter<Test>>::mutate(|v| {
+			*v = v.wrapping_add(1);
+			*v
+		});
+
 		// TODO: see https://github.com/paritytech/substrate/issues/2325
 		let mut res = vec![];
 		res.extend_from_slice(well_known_keys::CHILD_STORAGE_KEY_PREFIX);
 		res.extend_from_slice(b"default:");
-		res.extend_from_slice(&KEY_COUNTER.fetch_add(1, Ordering::Relaxed).to_le_bytes());
+		res.extend_from_slice(&new_seed.to_le_bytes());
 		res.extend_from_slice(&account_id.to_le_bytes());
 		res
 	}
@@ -146,6 +149,7 @@ impl ComputeDispatchFee<Call, u64> for DummyComputeDispatchFee {
 const ALICE: u64 = 1;
 const BOB: u64 = 2;
 const CHARLIE: u64 = 3;
+const DJANGO: u64 = 4;
 
 pub struct ExtBuilder {
 	existential_deposit: u64,
@@ -263,6 +267,7 @@ fn account_removal_removes_storage() {
 					deduct_block: System::block_number(),
 					code_hash: H256::repeat_byte(1),
 					rent_allowance: 40,
+					last_write: None,
 				}));
 
 				let mut overlay = OverlayAccountDb::<Test>::new(&DirectAccountDb);
@@ -277,6 +282,7 @@ fn account_removal_removes_storage() {
 					deduct_block: System::block_number(),
 					code_hash: H256::repeat_byte(2),
 					rent_allowance: 40,
+					last_write: None,
 				}));
 
 				let mut overlay = OverlayAccountDb::<Test>::new(&DirectAccountDb);
@@ -424,7 +430,7 @@ const HASH_DISPATCH_CALL: [u8; 32] = hex!("49dfdcaf9c1553be10634467e95b8e71a3bc1
 fn dispatch_call() {
 	// This test can fail due to the encoding changes. In case it becomes too annoying
 	// let's rewrite so as we use this module controlled call or we serialize it in runtime.
-	let encoded = parity_codec::Encode::encode(&Call::Balances(balances::Call::transfer(CHARLIE, 50)));
+	let encoded = Encode::encode(&Call::Balances(balances::Call::transfer(CHARLIE, 50)));
 	assert_eq!(&encoded[..], &hex!("00000300000000000000C8")[..]);
 
 	let wasm = wabt::wat2wasm(CODE_DISPATCH_CALL).unwrap();
@@ -555,7 +561,7 @@ const CODE_SET_RENT: &str = r#"
 	;; transfer 50 to ALICE
 	(func $call_2
 		(call $ext_dispatch_call
-			(i32.const 8)
+			(i32.const 68)
 			(i32.const 11)
 		)
 	)
@@ -627,10 +633,12 @@ const CODE_SET_RENT: &str = r#"
 	(data (i32.const 0) "\28")
 
 	;; Encoding of call transfer 50 to CHARLIE
-	(data (i32.const 8) "\00\00\03\00\00\00\00\00\00\00\C8")
+	(data (i32.const 68) "\00\00\03\00\00\00\00\00\00\00\C8")
 )
 "#;
-const HASH_SET_RENT: [u8; 32] = hex!("a51c2a6f3f68936d4ae9abdb93b28eedcbd0f6f39770e168f9025f0c1e7094ef");
+
+// Use test_hash_and_code test to get the actual hash if the code changed.
+const HASH_SET_RENT: [u8; 32] = hex!("21d6b1d59aa6038fcad632488e9026893a1bbb48581774c771b8f24320697f05");
 
 /// Input data for each call in set_rent code
 mod call {
@@ -643,10 +651,10 @@ mod call {
 /// Test correspondance of set_rent code and its hash.
 /// Also test that encoded extrinsic in code correspond to the correct transfer
 #[test]
-fn set_rent_hash_and_code() {
+fn test_set_rent_code_and_hash() {
 	// This test can fail due to the encoding changes. In case it becomes too annoying
 	// let's rewrite so as we use this module controlled call or we serialize it in runtime.
-	let encoded = parity_codec::Encode::encode(&Call::Balances(balances::Call::transfer(CHARLIE, 50)));
+	let encoded = Encode::encode(&Call::Balances(balances::Call::transfer(CHARLIE, 50)));
 	assert_eq!(&encoded[..], &hex!("00000300000000000000C8")[..]);
 
 	let wasm = wabt::wat2wasm(CODE_SET_RENT).unwrap();
@@ -657,7 +665,8 @@ fn set_rent_hash_and_code() {
 			Balances::deposit_creating(&ALICE, 1_000_000);
 			assert_ok!(Contract::put_code(Origin::signed(ALICE), 100_000, wasm));
 
-			// If you ever need to update the wasm source this test will fail and will show you the actual hash.
+			// If you ever need to update the wasm source this test will fail
+			// and will show you the actual hash.
 			assert_eq!(System::events(), vec![
 				EventRecord {
 					phase: Phase::ApplyExtrinsic(0),
@@ -691,15 +700,15 @@ fn storage_size() {
 				100_000, HASH_SET_RENT.into(),
 				<Test as balances::Trait>::Balance::from(1_000u32).encode() // rent allowance
 			));
-			let bob_contract = super::ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap();
+			let bob_contract = ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap();
 			assert_eq!(bob_contract.storage_size, Contract::storage_size_offset() + 4);
 
 			assert_ok!(Contract::call(Origin::signed(ALICE), BOB, 0, 100_000, call::set_storage_4_byte()));
-			let bob_contract = super::ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap();
+			let bob_contract = ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap();
 			assert_eq!(bob_contract.storage_size, Contract::storage_size_offset() + 4 + 4);
 
 			assert_ok!(Contract::call(Origin::signed(ALICE), BOB, 0, 100_000, call::remove_storage_4_byte()));
-			let bob_contract = super::ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap();
+			let bob_contract = ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap();
 			assert_eq!(bob_contract.storage_size, Contract::storage_size_offset() + 4);
 		}
 	);
@@ -723,7 +732,7 @@ fn deduct_blocks() {
 			));
 
 			// Check creation
-			let bob_contract = super::ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap();
+			let bob_contract = ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap();
 			assert_eq!(bob_contract.rent_allowance, 1_000);
 
 			// Advance 4 blocks
@@ -736,7 +745,7 @@ fn deduct_blocks() {
 			let rent = (8 + 4 - 3) // storage size = size_offset + deploy_set_storage - deposit_offset
 				* 4 // rent byte price
 				* 4; // blocks to rent
-			let bob_contract = super::ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap();
+			let bob_contract = ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap();
 			assert_eq!(bob_contract.rent_allowance, 1_000 - rent);
 			assert_eq!(bob_contract.deduct_block, 5);
 			assert_eq!(Balances::free_balance(BOB), 30_000 - rent);
@@ -751,7 +760,7 @@ fn deduct_blocks() {
 			let rent_2 = (8 + 4 - 2) // storage size = size_offset + deploy_set_storage - deposit_offset
 				* 4 // rent byte price
 				* 7; // blocks to rent
-			let bob_contract = super::ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap();
+			let bob_contract = ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap();
 			assert_eq!(bob_contract.rent_allowance, 1_000 - rent - rent_2);
 			assert_eq!(bob_contract.deduct_block, 12);
 			assert_eq!(Balances::free_balance(BOB), 30_000 - rent - rent_2);
@@ -759,7 +768,7 @@ fn deduct_blocks() {
 			// Second call on same block should have no effect on rent
 			assert_ok!(Contract::call(Origin::signed(ALICE), BOB, 0, 100_000, call::null()));
 
-			let bob_contract = super::ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap();
+			let bob_contract = ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap();
 			assert_eq!(bob_contract.rent_allowance, 1_000 - rent - rent_2);
 			assert_eq!(bob_contract.deduct_block, 12);
 			assert_eq!(Balances::free_balance(BOB), 30_000 - rent - rent_2);
@@ -822,9 +831,9 @@ fn claim_surcharge(blocks: u64, trigger_call: impl Fn() -> bool, removes: bool)
 			assert!(trigger_call());
 
 			if removes {
-				assert!(super::ContractInfoOf::<Test>::get(BOB).unwrap().get_tombstone().is_some());
+				assert!(ContractInfoOf::<Test>::get(BOB).unwrap().get_tombstone().is_some());
 			} else {
-				assert!(super::ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().is_some());
+				assert!(ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().is_some());
 			}
 		}
 	);
@@ -853,21 +862,21 @@ fn removals(trigger_call: impl Fn() -> bool) {
 
 			// Trigger rent must have no effect
 			assert!(trigger_call());
-			assert_eq!(super::ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap().rent_allowance, 1_000);
+			assert_eq!(ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap().rent_allowance, 1_000);
 
 			// Advance blocks
 			System::initialize(&10, &[0u8; 32].into(), &[0u8; 32].into());
 
 			// Trigger rent through call
 			assert!(trigger_call());
-			assert!(super::ContractInfoOf::<Test>::get(BOB).unwrap().get_tombstone().is_some());
+			assert!(ContractInfoOf::<Test>::get(BOB).unwrap().get_tombstone().is_some());
 
 			// Advance blocks
 			System::initialize(&20, &[0u8; 32].into(), &[0u8; 32].into());
 
 			// Trigger rent must have no effect
 			assert!(trigger_call());
-			assert!(super::ContractInfoOf::<Test>::get(BOB).unwrap().get_tombstone().is_some());
+			assert!(ContractInfoOf::<Test>::get(BOB).unwrap().get_tombstone().is_some());
 		}
 	);
 
@@ -887,21 +896,21 @@ fn removals(trigger_call: impl Fn() -> bool) {
 
 			// Trigger rent must have no effect
 			assert!(trigger_call());
-			assert_eq!(super::ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap().rent_allowance, 100);
+			assert_eq!(ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap().rent_allowance, 100);
 
 			// Advance blocks
 			System::initialize(&10, &[0u8; 32].into(), &[0u8; 32].into());
 
 			// Trigger rent through call
 			assert!(trigger_call());
-			assert!(super::ContractInfoOf::<Test>::get(BOB).unwrap().get_tombstone().is_some());
+			assert!(ContractInfoOf::<Test>::get(BOB).unwrap().get_tombstone().is_some());
 
 			// Advance blocks
 			System::initialize(&20, &[0u8; 32].into(), &[0u8; 32].into());
 
 			// Trigger rent must have no effect
 			assert!(trigger_call());
-			assert!(super::ContractInfoOf::<Test>::get(BOB).unwrap().get_tombstone().is_some());
+			assert!(ContractInfoOf::<Test>::get(BOB).unwrap().get_tombstone().is_some());
 		}
 	);
 
@@ -921,25 +930,25 @@ fn removals(trigger_call: impl Fn() -> bool) {
 
 			// Trigger rent must have no effect
 			assert!(trigger_call());
-			assert_eq!(super::ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap().rent_allowance, 1_000);
+			assert_eq!(ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap().rent_allowance, 1_000);
 
 			// Transfer funds
 			assert_ok!(Contract::call(Origin::signed(ALICE), BOB, 0, 100_000, call::transfer()));
-			assert_eq!(super::ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap().rent_allowance, 1_000);
+			assert_eq!(ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap().rent_allowance, 1_000);
 
 			// Advance blocks
 			System::initialize(&10, &[0u8; 32].into(), &[0u8; 32].into());
 
 			// Trigger rent through call
 			assert!(trigger_call());
-			assert!(super::ContractInfoOf::<Test>::get(BOB).is_none());
+			assert!(ContractInfoOf::<Test>::get(BOB).is_none());
 
 			// Advance blocks
 			System::initialize(&20, &[0u8; 32].into(), &[0u8; 32].into());
 
 			// Trigger rent must have no effect
 			assert!(trigger_call());
-			assert!(super::ContractInfoOf::<Test>::get(BOB).is_none());
+			assert!(ContractInfoOf::<Test>::get(BOB).is_none());
 		}
 	);
 }
@@ -1014,7 +1023,7 @@ fn default_rent_allowance_on_create() {
 			));
 
 			// Check creation
-			let bob_contract = super::ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap();
+			let bob_contract = ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap();
 			assert_eq!(bob_contract.rent_allowance, <BalanceOf<Test>>::max_value());
 
 			// Advance blocks
@@ -1024,8 +1033,196 @@ fn default_rent_allowance_on_create() {
 			assert_ok!(Contract::call(Origin::signed(ALICE), BOB, 0, 100_000, call::null()));
 
 			// Check contract is still alive
-			let bob_contract = super::ContractInfoOf::<Test>::get(BOB).unwrap().get_alive();
+			let bob_contract = ContractInfoOf::<Test>::get(BOB).unwrap().get_alive();
 			assert!(bob_contract.is_some())
 		}
 	);
 }
+
+const CODE_RESTORATION: &str = r#"
+(module
+	(import "env" "ext_set_storage" (func $ext_set_storage (param i32 i32 i32 i32)))
+	(import "env" "ext_dispatch_call" (func $ext_dispatch_call (param i32 i32)))
+	(import "env" "memory" (memory 1 1))
+
+	(func (export "call")
+		(call $ext_dispatch_call
+			(i32.const 200) ;; Pointer to the start of encoded call buffer
+			(i32.const 115) ;; Length of the buffer
+		)
+	)
+	(func (export "deploy")
+		;; Data to restore
+		(call $ext_set_storage
+			(i32.const 0)
+			(i32.const 1)
+			(i32.const 0)
+			(i32.const 4)
+		)
+
+		;; ACL
+		(call $ext_set_storage
+			(i32.const 100)
+			(i32.const 1)
+			(i32.const 0)
+			(i32.const 4)
+		)
+	)
+
+	;; Data to restore
+	(data (i32.const 0) "\28")
+
+	;; ACL
+	(data (i32.const 100) "\01")
+
+	;; Call
+	(data (i32.const 200) "\01\05\02\00\00\00\00\00\00\00\21\d6\b1\d5\9a\a6\03\8f\ca\d6\32\48\8e\90"
+		"\26\89\3a\1b\bb\48\58\17\74\c7\71\b8\f2\43\20\69\7f\05\32\00\00\00\00\00\00\00\08\01\00\00"
+		"\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\01"
+		"\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00"
+		"\00"
+	)
+)
+"#;
+const HASH_RESTORATION: [u8; 32] = hex!("b393bfa8de97b02f08ba1580b46100a80c9489a97c8d870691c9ff7236b29bc7");
+
+#[test]
+fn restorations_dirty_storage_and_different_storage() {
+	restoration(true, true);
+}
+
+#[test]
+fn restorations_dirty_storage() {
+	restoration(false, true);
+}
+
+#[test]
+fn restoration_different_storage() {
+	restoration(true, false);
+}
+
+#[test]
+fn restoration_success() {
+	restoration(false, false);
+}
+
+fn restoration(test_different_storage: bool, test_restore_to_with_dirty_storage: bool) {
+	let acl_key = {
+		let mut s = [0u8; 32];
+		s[0] = 1;
+		s
+	};
+
+	// This test can fail due to the encoding changes. In case it becomes too annoying
+	// let's rewrite so as we use this module controlled call or we serialize it in runtime.
+	let encoded = hex::encode(Encode::encode(&Call::Contract(super::Call::restore_to(
+		BOB,
+		HASH_SET_RENT.into(),
+		<Test as balances::Trait>::Balance::from(50u32),
+		vec![acl_key, acl_key],
+	))));
+
+	let literal = "0105020000000000000021d6b1d59aa6038fcad632488e9026893a1bbb48581774c771b8f243206\
+		97f053200000000000000080100000000000000000000000000000000000000000000000000000000000000010\
+		0000000000000000000000000000000000000000000000000000000000000";
+
+	assert_eq!(encoded, literal);
+	assert_eq!(115, hex::decode(literal).unwrap().len());
+
+	let restoration_wasm = wabt::wat2wasm(CODE_RESTORATION).unwrap();
+	let set_rent_wasm = wabt::wat2wasm(CODE_SET_RENT).unwrap();
+
+	with_externalities(
+		&mut ExtBuilder::default().existential_deposit(50).build(),
+		|| {
+			Balances::deposit_creating(&ALICE, 1_000_000);
+			assert_ok!(Contract::put_code(Origin::signed(ALICE), 100_000, restoration_wasm));
+			assert_ok!(Contract::put_code(Origin::signed(ALICE), 100_000, set_rent_wasm));
+
+			// If you ever need to update the wasm source this test will fail
+			// and will show you the actual hash.
+			assert_eq!(System::events(), vec![
+				EventRecord {
+					phase: Phase::ApplyExtrinsic(0),
+					event: MetaEvent::balances(balances::RawEvent::NewAccount(1, 1_000_000)),
+					topics: vec![],
+				},
+				EventRecord {
+					phase: Phase::ApplyExtrinsic(0),
+					event: MetaEvent::contract(RawEvent::CodeStored(HASH_RESTORATION.into())),
+					topics: vec![],
+				},
+				EventRecord {
+					phase: Phase::ApplyExtrinsic(0),
+					event: MetaEvent::contract(RawEvent::CodeStored(HASH_SET_RENT.into())),
+					topics: vec![],
+				},
+			]);
+
+			assert_ok!(Contract::create(
+				Origin::signed(ALICE),
+				30_000,
+				100_000, HASH_SET_RENT.into(),
+				<Test as balances::Trait>::Balance::from(0u32).encode()
+			));
+
+			// Check creation
+			let bob_contract = ContractInfoOf::<Test>::get(BOB).unwrap().get_alive().unwrap();
+			assert_eq!(bob_contract.rent_allowance, 0);
+
+			if test_different_storage {
+				assert_ok!(Contract::call(
+					Origin::signed(ALICE),
+					BOB, 0, 100_000,
+					call::set_storage_4_byte())
+				);
+			}
+
+			// Advance 4 blocks
+			System::initialize(&5, &[0u8; 32].into(), &[0u8; 32].into());
+
+			// Trigger rent through call
+			assert_ok!(Contract::call(Origin::signed(ALICE), BOB, 0, 100_000, call::null()));
+			assert!(ContractInfoOf::<Test>::get(BOB).unwrap().get_tombstone().is_some());
+
+			Balances::deposit_creating(&CHARLIE, 1_000_000);
+			assert_ok!(Contract::create(
+				Origin::signed(CHARLIE),
+				30_000,
+				100_000, HASH_RESTORATION.into(),
+				<Test as balances::Trait>::Balance::from(0u32).encode()
+			));
+
+			let django_trie_id = ContractInfoOf::<Test>::get(DJANGO).unwrap()
+				.get_alive().unwrap().trie_id;
+
+			if !test_restore_to_with_dirty_storage {
+				// Advance 1 blocks
+				System::initialize(&6, &[0u8; 32].into(), &[0u8; 32].into());
+			}
+
+			assert_ok!(Contract::call(
+				Origin::signed(ALICE),
+				DJANGO, 0, 100_000,
+				vec![],
+			));
+
+			if test_different_storage || test_restore_to_with_dirty_storage {
+				assert!(ContractInfoOf::<Test>::get(BOB).unwrap().get_tombstone().is_some());
+				let django_contract = ContractInfoOf::<Test>::get(DJANGO).unwrap()
+					.get_alive().unwrap();
+				assert_eq!(django_contract.storage_size, 16);
+				assert_eq!(django_contract.trie_id, django_trie_id);
+				assert_eq!(django_contract.deduct_block, System::block_number());
+			} else {
+				let bob_contract = ContractInfoOf::<Test>::get(BOB).unwrap()
+					.get_alive().unwrap();
+				assert_eq!(bob_contract.rent_allowance, 50);
+				assert_eq!(bob_contract.storage_size, 12);
+				assert_eq!(bob_contract.trie_id, django_trie_id);
+				assert_eq!(bob_contract.deduct_block, System::block_number());
+				assert!(ContractInfoOf::<Test>::get(DJANGO).is_none());
+			}
+		}
+	);
+}

substrate/srml/executive/src/lib.rs

@@ -474,7 +474,7 @@ mod tests {
 				header: Header {
 					parent_hash: [69u8; 32].into(),
 					number: 1,
-					state_root: hex!("ac2840371d51ff2e036c8fc05af7313b7a030f735c38b2f03b94cbe87bfbb7c9").into(),
+					state_root: hex!("5ba497e45e379d80a4524f9509d224e9c175d0fa30f3491481e7e44a6a758adf").into(),
 					extrinsics_root: hex!("03170a2e7597b7b7e3d84c05391d139a62b157e78786d8c082f29dcf4c111314").into(),
 					digest: Digest { logs: vec![], },
 				},