contracts: Charge rent for code storage (#7935)

* contracts: Implement refcounting for wasm code

* contracts: Charge rent for code storage

* contracts: Fix dispatchables erroneously refunding base costs

* Fixed typos in comments.

Co-authored-by: Andrew Jones <ascjones@gmail.com>

* Remove awkward empty line

* Fix more typos in docs

* Fix typos in docs

Co-authored-by: Andrew Jones <ascjones@gmail.com>

* Split up complicated expression

Co-authored-by: Andrew Jones <ascjones@gmail.com>

* review: Remove unused return value

* Fix typos

Co-authored-by: Andrew Jones <ascjones@gmail.com>

* review: Fix refcount being reset to one on re-instrumentation

* Document evictable_code parameter

* Make Executable::execute consume and store itself

* Added comments about stale values

* Disregard struct size in occupied_storage()

Co-authored-by: Andrew Jones <ascjones@gmail.com>

substrate/frame/contracts/src/wasm/code_cache.rs

@@ -27,46 +27,84 @@
 //! this guarantees that every instrumented contract code in cache cannot have the version equal to the current one.
 //! Thus, before executing a contract it should be reinstrument with new schedule.
 
-use crate::wasm::{prepare, runtime::Env, PrefabWasmModule};
-use crate::{CodeHash, CodeStorage, PristineCode, Schedule, Config};
-use sp_std::prelude::*;
-use sp_runtime::traits::Hash;
+use crate::{
+	CodeHash, CodeStorage, PristineCode, Schedule, Config, Error,
+	wasm::{prepare, PrefabWasmModule}, Module as Contracts, RawEvent,
+};
 use sp_core::crypto::UncheckedFrom;
-use frame_support::StorageMap;
+use frame_support::{StorageMap, dispatch::{DispatchError, DispatchResult}};
 
-/// Put code in the storage. The hash of code is used as a key and is returned
-/// as a result of this function.
+/// Put the instrumented module in storage.
 ///
-/// This function instruments the given code and caches it in the storage.
-pub fn save<T: Config>(
-	original_code: Vec<u8>,
-	schedule: &Schedule<T>,
-) -> Result<CodeHash<T>, &'static str> where T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]> {
-	let prefab_module = prepare::prepare_contract::<Env, T>(&original_code, schedule)?;
-	let code_hash = T::Hashing::hash(&original_code);
+/// Increments the refcount of the in-storage `prefab_module` if it already exists in storage
+/// under the specified `code_hash`.
+pub fn store<T: Config>(mut prefab_module: PrefabWasmModule<T>)
+where
+	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>
+{
+	let code_hash = sp_std::mem::take(&mut prefab_module.code_hash);
 
-	<CodeStorage<T>>::insert(code_hash, prefab_module);
-	<PristineCode<T>>::insert(code_hash, original_code);
-
-	Ok(code_hash)
+	// original_code is only `Some` if the contract was instantiated from a new code
+	// but `None` if it was loaded from storage.
+	if let Some(code) = prefab_module.original_code.take() {
+		<PristineCode<T>>::insert(&code_hash, code);
+	}
+	<CodeStorage<T>>::mutate(&code_hash, |existing| {
+		match existing {
+			Some(module) => increment_64(&mut module.refcount),
+			None => {
+				*existing = Some(prefab_module);
+				Contracts::<T>::deposit_event(RawEvent::CodeStored(code_hash))
+			}
+		}
+	});
 }
 
-/// Version of `save` to be used in runtime benchmarks.
-//
-/// This version neither checks nor instruments the passed in code. This is useful
-/// when code needs to be benchmarked without the injected instrumentation.
-#[cfg(feature = "runtime-benchmarks")]
-pub fn save_raw<T: Config>(
-	original_code: Vec<u8>,
-	schedule: &Schedule<T>,
-) -> Result<CodeHash<T>, &'static str> where T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]> {
-	let prefab_module = prepare::benchmarking::prepare_contract::<T>(&original_code, schedule)?;
-	let code_hash = T::Hashing::hash(&original_code);
+/// Decrement the refcount and store.
+///
+/// Removes the code instead of storing it when the refcount drops to zero.
+pub fn store_decremented<T: Config>(mut prefab_module: PrefabWasmModule<T>)
+where
+	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>
+{
+	prefab_module.refcount = prefab_module.refcount.saturating_sub(1);
+	if prefab_module.refcount > 0 {
+		<CodeStorage<T>>::insert(prefab_module.code_hash, prefab_module);
+	} else {
+		<CodeStorage<T>>::remove(prefab_module.code_hash);
+		finish_removal::<T>(prefab_module.code_hash);
+	}
+}
 
-	<CodeStorage<T>>::insert(code_hash, prefab_module);
-	<PristineCode<T>>::insert(code_hash, original_code);
+/// Increment the refcount of a code in-storage by one.
+pub fn increment_refcount<T: Config>(code_hash: CodeHash<T>) -> DispatchResult
+where
+	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>
+{
+	<CodeStorage<T>>::mutate(code_hash, |existing| {
+		if let Some(module) = existing {
+			increment_64(&mut module.refcount);
+			Ok(())
+		} else {
+			Err(Error::<T>::CodeNotFound.into())
+		}
+	})
+}
 
-	Ok(code_hash)
+/// Decrement the refcount of a code in-storage by one and remove the code when it drops to zero.
+pub fn decrement_refcount<T: Config>(code_hash: CodeHash<T>)
+where
+	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>
+{
+	<CodeStorage<T>>::mutate_exists(code_hash, |existing| {
+		if let Some(module) = existing {
+			module.refcount = module.refcount.saturating_sub(1);
+			if module.refcount == 0 {
+				*existing = None;
+				finish_removal::<T>(code_hash);
+			}
+		}
+	});
 }
 
 /// Load code with the given code hash.
@@ -75,21 +113,51 @@ pub fn save_raw<T: Config>(
 /// the current one given as an argument, then this function will perform
 /// re-instrumentation and update the cache in the storage.
 pub fn load<T: Config>(
-	code_hash: &CodeHash<T>,
-	schedule: &Schedule<T>,
-) -> Result<PrefabWasmModule, &'static str> where T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]> {
-	let mut prefab_module =
-		<CodeStorage<T>>::get(code_hash).ok_or_else(|| "code is not found")?;
+	code_hash: CodeHash<T>,
+	schedule: Option<&Schedule<T>>,
+) -> Result<PrefabWasmModule<T>, DispatchError>
+where
+	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>
+{
+	let mut prefab_module = <CodeStorage<T>>::get(code_hash)
+		.ok_or_else(|| Error::<T>::CodeNotFound)?;
 
-	if prefab_module.schedule_version < schedule.version {
-		// The current schedule version is greater than the version of the one cached
-		// in the storage.
-		//
-		// We need to re-instrument the code with the latest schedule here.
-		let original_code =
-			<PristineCode<T>>::get(code_hash).ok_or_else(|| "pristine code is not found")?;
-		prefab_module = prepare::prepare_contract::<Env, T>(&original_code, schedule)?;
-		<CodeStorage<T>>::insert(&code_hash, &prefab_module);
+	if let Some(schedule) = schedule {
+		if prefab_module.schedule_version < schedule.version {
+			// The current schedule version is greater than the version of the one cached
+			// in the storage.
+			//
+			// We need to re-instrument the code with the latest schedule here.
+			let original_code = <PristineCode<T>>::get(code_hash)
+				.ok_or_else(|| Error::<T>::CodeNotFound)?;
+			prefab_module.code = prepare::reinstrument_contract::<T>(original_code, schedule)?;
+			prefab_module.schedule_version = schedule.version;
+			<CodeStorage<T>>::insert(&code_hash, &prefab_module);
+		}
 	}
+	prefab_module.code_hash = code_hash;
 	Ok(prefab_module)
 }
+
+/// Finish removal of a code by deleting the pristine code and emitting an event.
+fn finish_removal<T: Config>(code_hash: CodeHash<T>)
+where
+	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>
+{
+	<PristineCode<T>>::remove(code_hash);
+	Contracts::<T>::deposit_event(RawEvent::CodeRemoved(code_hash))
+}
+
+/// Increment the refcount panicking if it should ever overflow (which will not happen).
+///
+/// We try hard to be infallible here because otherwise more storage transactions would be
+/// necessary to account for failures in storing code for an already instantiated contract.
+fn increment_64(refcount: &mut u64) {
+	*refcount = refcount.checked_add(1).expect("
+		refcount is 64bit. Generating this overflow would require to store
+		_at least_ 18 exabyte of data assuming that a contract consumes only
+		one byte of data. Any node would run out of storage space before hitting
+		this overflow.
+		qed
+	");
+}

substrate/frame/contracts/src/wasm/mod.rs

@@ -18,114 +18,163 @@
 //! This module provides a means for executing contracts
 //! represented in wasm.
 
-use crate::{
-	CodeHash, Schedule, Config,
-	wasm::env_def::FunctionImplProvider,
-	exec::Ext,
-	gas::GasMeter,
-};
-use sp_std::prelude::*;
-use sp_core::crypto::UncheckedFrom;
-use codec::{Encode, Decode};
-
 #[macro_use]
 mod env_def;
 mod code_cache;
 mod prepare;
 mod runtime;
 
-use self::code_cache::load as load_code;
+use crate::{
+	CodeHash, Schedule, Config,
+	wasm::env_def::FunctionImplProvider,
+	exec::{Ext, Executable, ExportedFunction},
+	gas::GasMeter,
+};
+use sp_std::prelude::*;
+use sp_core::crypto::UncheckedFrom;
+use codec::{Encode, Decode};
+use frame_support::dispatch::{DispatchError, DispatchResult};
 use pallet_contracts_primitives::ExecResult;
-
-pub use self::code_cache::save as save_code;
-#[cfg(feature = "runtime-benchmarks")]
-pub use self::code_cache::save_raw as save_code_raw;
 pub use self::runtime::{ReturnCode, Runtime, RuntimeToken};
 
 /// A prepared wasm module ready for execution.
+///
+/// # Note
+///
+/// This data structure is mostly immutable once created and stored. The exceptions that
+/// can be changed by calling a contract are `refcount`, `schedule_version` and `code`.
+/// `refcount` can change when a contract instantiates a new contract or self terminates.
+/// `schedule_version` and `code` when a contract with an outdated instrumention is called.
+/// Therefore one must be careful when holding any in-memory representation of this type while
+/// calling into a contract as those fields can get out of date.
 #[derive(Clone, Encode, Decode)]
-pub struct PrefabWasmModule {
+pub struct PrefabWasmModule<T: Config> {
 	/// Version of the schedule with which the code was instrumented.
 	#[codec(compact)]
 	schedule_version: u32,
+	/// Initial memory size of a contract's sandbox.
 	#[codec(compact)]
 	initial: u32,
+	/// The maximum memory size of a contract's sandbox.
 	#[codec(compact)]
 	maximum: u32,
+	/// The number of alive contracts that use this as their contract code.
+	///
+	/// If this number drops to zero this module is removed from storage.
+	#[codec(compact)]
+	refcount: u64,
 	/// This field is reserved for future evolution of format.
 	///
-	/// Basically, for now this field will be serialized as `None`. In the future
-	/// we would be able to extend this structure with.
+	/// For now this field is serialized as `None`. In the future we are able to change the
+	/// type parameter to a new struct that contains the fields that we want to add.
+	/// That new struct would also contain a reserved field for its future extensions.
+	/// This works because in SCALE `None` is encoded independently from the type parameter
+	/// of the option.
 	_reserved: Option<()>,
 	/// Code instrumented with the latest schedule.
 	code: Vec<u8>,
+	/// The size of the uninstrumented code.
+	///
+	/// We cache this value here in order to avoid the need to pull the pristine code
+	/// from storage when we only need its length for rent calculations.
+	original_code_len: u32,
+	/// The uninstrumented, pristine version of the code.
+	///
+	/// It is not stored because the pristine code has its own storage item. The value
+	/// is only `Some` when this module was created from an `original_code` and `None` if
+	/// it was loaded from storage.
+	#[codec(skip)]
+	original_code: Option<Vec<u8>>,
+	/// The code hash of the stored code which is defined as the hash over the `original_code`.
+	///
+	/// As the map key there is no need to store the hash in the value, too. It is set manually
+	/// when loading the module from storage.
+	#[codec(skip)]
+	code_hash: CodeHash<T>,
 }
 
-/// Wasm executable loaded by `WasmLoader` and executed by `WasmVm`.
-pub struct WasmExecutable {
-	entrypoint_name: &'static str,
-	prefab_module: PrefabWasmModule,
-}
-
-/// Loader which fetches `WasmExecutable` from the code cache.
-pub struct WasmLoader<'a, T: Config> {
-	schedule: &'a Schedule<T>,
-}
-
-impl<'a, T: Config> WasmLoader<'a, T> where T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]> {
-	pub fn new(schedule: &'a Schedule<T>) -> Self {
-		WasmLoader { schedule }
+impl ExportedFunction {
+	/// The wasm export name for the function.
+	fn identifier(&self) -> &str {
+		match self {
+			Self::Constructor => "deploy",
+			Self::Call => "call",
+		}
 	}
 }
 
-impl<'a, T: Config> crate::exec::Loader<T> for WasmLoader<'a, T>
+impl<T: Config> PrefabWasmModule<T>
 where
 	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>
 {
-	type Executable = WasmExecutable;
-
-	fn load_init(&self, code_hash: &CodeHash<T>) -> Result<WasmExecutable, &'static str> {
-		let prefab_module = load_code::<T>(code_hash, self.schedule)?;
-		Ok(WasmExecutable {
-			entrypoint_name: "deploy",
-			prefab_module,
-		})
+	/// Create the module by checking and instrumenting `original_code`.
+	pub fn from_code(
+		original_code: Vec<u8>,
+		schedule: &Schedule<T>
+	) -> Result<Self, DispatchError> {
+		prepare::prepare_contract(original_code, schedule).map_err(Into::into)
 	}
-	fn load_main(&self, code_hash: &CodeHash<T>) -> Result<WasmExecutable, &'static str> {
-		let prefab_module = load_code::<T>(code_hash, self.schedule)?;
-		Ok(WasmExecutable {
-			entrypoint_name: "call",
-			prefab_module,
-		})
+
+	/// Create and store the module without checking nor instrumenting the passed code.
+	///
+	/// # Note
+	///
+	/// This is useful for benchmarking where we don't want instrumentation to skew
+	/// our results.
+	#[cfg(feature = "runtime-benchmarks")]
+	pub fn store_code_unchecked(
+		original_code: Vec<u8>,
+		schedule: &Schedule<T>
+	) -> DispatchResult {
+		let executable = prepare::benchmarking::prepare_contract(original_code, schedule)
+			.map_err::<DispatchError, _>(Into::into)?;
+		code_cache::store(executable);
+		Ok(())
+	}
+
+	/// Return the refcount of the module.
+	#[cfg(test)]
+	pub fn refcount(&self) -> u64 {
+		self.refcount
 	}
 }
 
-/// Implementation of `Vm` that takes `WasmExecutable` and executes it.
-pub struct WasmVm<'a, T: Config> where T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]> {
-	schedule: &'a Schedule<T>,
-}
-
-impl<'a, T: Config> WasmVm<'a, T> where T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]> {
-	pub fn new(schedule: &'a Schedule<T>) -> Self {
-		WasmVm { schedule }
-	}
-}
-
-impl<'a, T: Config> crate::exec::Vm<T> for WasmVm<'a, T>
+impl<T: Config> Executable<T> for PrefabWasmModule<T>
 where
 	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>
 {
-	type Executable = WasmExecutable;
+	fn from_storage(
+		code_hash: CodeHash<T>,
+		schedule: &Schedule<T>
+	) -> Result<Self, DispatchError> {
+		code_cache::load(code_hash, Some(schedule))
+	}
+
+	fn from_storage_noinstr(code_hash: CodeHash<T>) -> Result<Self, DispatchError> {
+		code_cache::load(code_hash, None)
+	}
+
+	fn drop_from_storage(self) {
+		code_cache::store_decremented(self);
+	}
+
+	fn add_user(code_hash: CodeHash<T>) -> DispatchResult {
+		code_cache::increment_refcount::<T>(code_hash)
+	}
+
+	fn remove_user(code_hash: CodeHash<T>) {
+		code_cache::decrement_refcount::<T>(code_hash)
+	}
 
 	fn execute<E: Ext<T = T>>(
-		&self,
-		exec: &WasmExecutable,
+		self,
 		mut ext: E,
+		function: &ExportedFunction,
 		input_data: Vec<u8>,
 		gas_meter: &mut GasMeter<E::T>,
 	) -> ExecResult {
 		let memory =
-			sp_sandbox::Memory::new(exec.prefab_module.initial, Some(exec.prefab_module.maximum))
+			sp_sandbox::Memory::new(self.initial, Some(self.maximum))
 				.unwrap_or_else(|_| {
 				// unlike `.expect`, explicit panic preserves the source location.
 				// Needed as we can't use `RUST_BACKTRACE` in here.
@@ -145,17 +194,34 @@ where
 		let mut runtime = Runtime::new(
 			&mut ext,
 			input_data,
-			&self.schedule,
 			memory,
 			gas_meter,
 		);
 
+		// We store before executing so that the code hash is available in the constructor.
+		let code = self.code.clone();
+		if let &ExportedFunction::Constructor = function {
+			code_cache::store(self)
+		}
+
 		// Instantiate the instance from the instrumented module code and invoke the contract
 		// entrypoint.
-		let result = sp_sandbox::Instance::new(&exec.prefab_module.code, &imports, &mut runtime)
-			.and_then(|mut instance| instance.invoke(exec.entrypoint_name, &[], &mut runtime));
+		let result = sp_sandbox::Instance::new(&code, &imports, &mut runtime)
+			.and_then(|mut instance| instance.invoke(function.identifier(), &[], &mut runtime));
+
 		runtime.to_execution_result(result)
 	}
+
+	fn code_hash(&self) -> &CodeHash<T> {
+		&self.code_hash
+	}
+
+	fn occupied_storage(&self) -> u32 {
+		// We disregard the size of the struct itself as the size is completely
+		// dominated by the code size.
+		let len = self.original_code_len.saturating_add(self.code.len() as u32);
+		len.checked_div(self.refcount as u32).unwrap_or(len)
+	}
 }
 
 #[cfg(test)]
@@ -163,10 +229,9 @@ mod tests {
 	use super::*;
 	use crate::{
 		CodeHash, BalanceOf, Error, Module as Contracts,
-		exec::{Ext, StorageKey, AccountIdOf},
+		exec::{Ext, StorageKey, AccountIdOf, Executable},
 		gas::{Gas, GasMeter},
 		tests::{Test, Call, ALICE, BOB},
-		wasm::prepare::prepare_contract,
 	};
 	use std::collections::HashMap;
 	use sp_core::H256;
@@ -220,6 +285,7 @@ mod tests {
 		restores: Vec<RestoreEntry>,
 		// (topics, data)
 		events: Vec<(Vec<H256>, Vec<u8>)>,
+		schedule: Schedule<Test>,
 	}
 
 	impl Ext for MockExt {
@@ -234,7 +300,7 @@ mod tests {
 		}
 		fn instantiate(
 			&mut self,
-			code_hash: &CodeHash<Test>,
+			code_hash: CodeHash<Test>,
 			endowment: u64,
 			gas_meter: &mut GasMeter<Test>,
 			data: Vec<u8>,
@@ -248,7 +314,7 @@ mod tests {
 				salt: salt.to_vec(),
 			});
 			Ok((
-				Contracts::<Test>::contract_address(&ALICE, code_hash, salt),
+				Contracts::<Test>::contract_address(&ALICE, &code_hash, salt),
 				ExecReturnValue {
 					flags: ReturnFlags::empty(),
 					data: Vec::new(),
@@ -355,6 +421,10 @@ mod tests {
 		fn get_weight_price(&self, weight: Weight) -> BalanceOf<Self::T> {
 			BalanceOf::<Self::T>::from(1312_u32).saturating_mul(weight.into())
 		}
+
+		fn schedule(&self) -> &Schedule<Self::T> {
+			&self.schedule
+		}
 	}
 
 	impl Ext for &mut MockExt {
@@ -368,7 +438,7 @@ mod tests {
 		}
 		fn instantiate(
 			&mut self,
-			code: &CodeHash<Test>,
+			code: CodeHash<Test>,
 			value: u64,
 			gas_meter: &mut GasMeter<Test>,
 			input_data: Vec<u8>,
@@ -454,6 +524,9 @@ mod tests {
 		fn get_weight_price(&self, weight: Weight) -> BalanceOf<Self::T> {
 			(**self).get_weight_price(weight)
 		}
+		fn schedule(&self) -> &Schedule<Self::T> {
+			(**self).schedule()
+		}
 	}
 
 	fn execute<E: Ext>(
@@ -466,23 +539,10 @@ mod tests {
 		<E::T as frame_system::Config>::AccountId:
 			UncheckedFrom<<E::T as frame_system::Config>::Hash> + AsRef<[u8]>
 	{
-		use crate::exec::Vm;
-
 		let wasm = wat::parse_str(wat).unwrap();
 		let schedule = crate::Schedule::default();
-		let prefab_module =
-			prepare_contract::<super::runtime::Env, E::T>(&wasm, &schedule).unwrap();
-
-		let exec = WasmExecutable {
-			// Use a "call" convention.
-			entrypoint_name: "call",
-			prefab_module,
-		};
-
-		let cfg = Default::default();
-		let vm = WasmVm::new(&cfg);
-
-		vm.execute(&exec, ext, input_data, gas_meter)
+		let executable = PrefabWasmModule::<E::T>::from_code(wasm, &schedule).unwrap();
+		executable.execute(ext, &ExportedFunction::Call, input_data, gas_meter)
 	}
 
 	const CODE_TRANSFER: &str = r#"

substrate/frame/contracts/src/wasm/prepare.rs

@@ -25,7 +25,7 @@ use crate::{
 	wasm::{PrefabWasmModule, env_def::ImportSatisfyCheck},
 };
 use parity_wasm::elements::{self, Internal, External, MemoryType, Type, ValueType};
-use pwasm_utils;
+use sp_runtime::traits::Hash;
 use sp_std::prelude::*;
 
 /// Currently, all imported functions must be located inside this module. We might support
@@ -407,22 +407,11 @@ fn get_memory_limits<T: Config>(module: Option<&MemoryType>, schedule: &Schedule
 	}
 }
 
-/// Loads the given module given in `original_code`, performs some checks on it and
-/// does some preprocessing.
-///
-/// The checks are:
-///
-/// - provided code is a valid wasm module.
-/// - the module doesn't define an internal memory instance,
-/// - imported memory (if any) doesn't reserve more memory than permitted by the `schedule`,
-/// - all imported functions from the external environment matches defined by `env` module,
-///
-/// The preprocessing includes injecting code for gas metering and metering the height of stack.
-pub fn prepare_contract<C: ImportSatisfyCheck, T: Config>(
+fn check_and_instrument<C: ImportSatisfyCheck, T: Config>(
 	original_code: &[u8],
 	schedule: &Schedule<T>,
-) -> Result<PrefabWasmModule, &'static str> {
-	let mut contract_module = ContractModule::new(original_code, schedule)?;
+) -> Result<(Vec<u8>, (u32, u32)), &'static str> {
+	let contract_module = ContractModule::new(&original_code, schedule)?;
 	contract_module.scan_exports()?;
 	contract_module.ensure_no_internal_memory()?;
 	contract_module.ensure_table_size_limit(schedule.limits.table_size)?;
@@ -438,19 +427,65 @@ pub fn prepare_contract<C: ImportSatisfyCheck, T: Config>(
 		schedule
 	)?;
 
-	contract_module = contract_module
+	let code = contract_module
 		.inject_gas_metering()?
-		.inject_stack_height_metering()?;
+		.inject_stack_height_metering()?
+		.into_wasm_code()?;
 
+	Ok((code, memory_limits))
+}
+
+fn do_preparation<C: ImportSatisfyCheck, T: Config>(
+	original_code: Vec<u8>,
+	schedule: &Schedule<T>,
+) -> Result<PrefabWasmModule<T>, &'static str> {
+	let (code, (initial, maximum)) = check_and_instrument::<C, T>(
+		original_code.as_ref(),
+		schedule,
+	)?;
 	Ok(PrefabWasmModule {
 		schedule_version: schedule.version,
-		initial: memory_limits.0,
-		maximum: memory_limits.1,
+		initial,
+		maximum,
 		_reserved: None,
-		code: contract_module.into_wasm_code()?,
+		code,
+		original_code_len: original_code.len() as u32,
+		refcount: 1,
+		code_hash: T::Hashing::hash(&original_code),
+		original_code: Some(original_code),
 	})
 }
 
+/// Loads the given module given in `original_code`, performs some checks on it and
+/// does some preprocessing.
+///
+/// The checks are:
+///
+/// - provided code is a valid wasm module.
+/// - the module doesn't define an internal memory instance,
+/// - imported memory (if any) doesn't reserve more memory than permitted by the `schedule`,
+/// - all imported functions from the external environment matches defined by `env` module,
+///
+/// The preprocessing includes injecting code for gas metering and metering the height of stack.
+pub fn prepare_contract<T: Config>(
+	original_code: Vec<u8>,
+	schedule: &Schedule<T>,
+) -> Result<PrefabWasmModule<T>, &'static str> {
+	do_preparation::<super::runtime::Env, T>(original_code, schedule)
+}
+
+/// The same as [`prepare_contract`] but without constructing a new [`PrefabWasmModule`]
+///
+/// # Note
+///
+/// Use this when an existing contract should be re-instrumented with a newer schedule version.
+pub fn reinstrument_contract<T: Config>(
+	original_code: Vec<u8>,
+	schedule: &Schedule<T>,
+) -> Result<Vec<u8>, &'static str> {
+	Ok(check_and_instrument::<super::runtime::Env, T>(&original_code, schedule)?.0)
+}
+
 /// Alternate (possibly unsafe) preparation functions used only for benchmarking.
 ///
 /// For benchmarking we need to construct special contracts that might not pass our
@@ -459,9 +494,7 @@ pub fn prepare_contract<C: ImportSatisfyCheck, T: Config>(
 /// in production code.
 #[cfg(feature = "runtime-benchmarks")]
 pub mod benchmarking {
-	use super::{
-		Config, ContractModule, PrefabWasmModule, ImportSatisfyCheck, Schedule, get_memory_limits
-	};
+	use super::*;
 	use parity_wasm::elements::FunctionType;
 
 	impl ImportSatisfyCheck for () {
@@ -471,10 +504,10 @@ pub mod benchmarking {
 	}
 
 	/// Prepare function that neither checks nor instruments the passed in code.
-	pub fn prepare_contract<T: Config>(original_code: &[u8], schedule: &Schedule<T>)
-		-> Result<PrefabWasmModule, &'static str>
+	pub fn prepare_contract<T: Config>(original_code: Vec<u8>, schedule: &Schedule<T>)
+		-> Result<PrefabWasmModule<T>, &'static str>
 	{
-		let contract_module = ContractModule::new(original_code, schedule)?;
+		let contract_module = ContractModule::new(&original_code, schedule)?;
 		let memory_limits = get_memory_limits(contract_module.scan_imports::<()>(&[])?, schedule)?;
 		Ok(PrefabWasmModule {
 			schedule_version: schedule.version,
@@ -482,6 +515,10 @@ pub mod benchmarking {
 			maximum: memory_limits.1,
 			_reserved: None,
 			code: contract_module.into_wasm_code()?,
+			original_code_len: original_code.len() as u32,
+			refcount: 1,
+			code_hash: T::Hashing::hash(&original_code),
+			original_code: Some(original_code),
 		})
 	}
 }
@@ -493,7 +530,7 @@ mod tests {
 	use std::fmt;
 	use assert_matches::assert_matches;
 
-	impl fmt::Debug for PrefabWasmModule {
+	impl fmt::Debug for PrefabWasmModule<crate::tests::Test> {
 		fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 			write!(f, "PreparedContract {{ .. }}")
 		}
@@ -534,7 +571,7 @@ mod tests {
 					},
 					.. Default::default()
 				};
-				let r = prepare_contract::<env::Test, crate::tests::Test>(wasm.as_ref(), &schedule);
+				let r = do_preparation::<env::Test, crate::tests::Test>(wasm, &schedule);
 				assert_matches!(r, $($expected)*);
 			}
 		};
@@ -945,7 +982,7 @@ mod tests {
 			).unwrap();
 			let mut schedule = Schedule::default();
 			schedule.enable_println = true;
-			let r = prepare_contract::<env::Test, crate::tests::Test>(wasm.as_ref(), &schedule);
+			let r = do_preparation::<env::Test, crate::tests::Test>(wasm, &schedule);
 			assert_matches!(r, Ok(_));
 		}
 	}

substrate/frame/contracts/src/wasm/runtime.rs

@@ -18,7 +18,7 @@
 //! Environment definition of the wasm smart-contract runtime.
 
 use crate::{
-	HostFnWeights, Schedule, Config, CodeHash, BalanceOf, Error,
+	HostFnWeights, Config, CodeHash, BalanceOf, Error,
 	exec::{Ext, StorageKey, TopicOf},
 	gas::{Gas, GasMeter, Token, GasMeterResult, ChargedAmount},
 	wasm::env_def::ConvertibleToWasm,
@@ -300,7 +300,6 @@ fn has_duplicates<T: PartialEq + AsRef<[u8]>>(items: &mut Vec<T>) -> bool {
 pub struct Runtime<'a, E: Ext + 'a> {
 	ext: &'a mut E,
 	input_data: Option<Vec<u8>>,
-	schedule: &'a Schedule<E::T>,
 	memory: sp_sandbox::Memory,
 	gas_meter: &'a mut GasMeter<E::T>,
 	trap_reason: Option<TrapReason>,
@@ -315,14 +314,12 @@ where
 	pub fn new(
 		ext: &'a mut E,
 		input_data: Vec<u8>,
-		schedule: &'a Schedule<E::T>,
 		memory: sp_sandbox::Memory,
 		gas_meter: &'a mut GasMeter<E::T>,
 	) -> Self {
 		Runtime {
 			ext,
 			input_data: Some(input_data),
-			schedule,
 			memory,
 			gas_meter,
 			trap_reason: None,
@@ -411,7 +408,7 @@ where
 	where
 		Tok: Token<E::T, Metadata=HostFnWeights<E::T>>,
 	{
-		match self.gas_meter.charge(&self.schedule.host_fn_weights, token) {
+		match self.gas_meter.charge(&self.ext.schedule().host_fn_weights, token) {
 			GasMeterResult::Proceed(amount) => Ok(amount),
 			GasMeterResult::OutOfGas => Err(Error::<E::T>::OutOfGas.into())
 		}
@@ -425,7 +422,7 @@ where
 	pub fn read_sandbox_memory(&self, ptr: u32, len: u32)
 	-> Result<Vec<u8>, DispatchError>
 	{
-		ensure!(len <= self.schedule.limits.max_memory_size(), Error::<E::T>::OutOfBounds);
+		ensure!(len <= self.ext.schedule().limits.max_memory_size(), Error::<E::T>::OutOfBounds);
 		let mut buf = vec![0u8; len as usize];
 		self.memory.get(ptr, buf.as_mut_slice())
 			.map_err(|_| Error::<E::T>::OutOfBounds)?;
@@ -889,7 +886,7 @@ define_env!(Env, <E: Ext>,
 			match nested_meter {
 				Some(nested_meter) => {
 					ext.instantiate(
-						&code_hash,
+						code_hash,
 						value,
 						nested_meter,
 						input_data,
@@ -1094,7 +1091,7 @@ define_env!(Env, <E: Ext>,
 	// The data is encoded as T::Hash.
 	seal_random(ctx, subject_ptr: u32, subject_len: u32, out_ptr: u32, out_len_ptr: u32) => {
 		ctx.charge_gas(RuntimeToken::Random)?;
-		if subject_len > ctx.schedule.limits.subject_len {
+		if subject_len > ctx.ext.schedule().limits.subject_len {
 			Err(Error::<E::T>::RandomSubjectTooLong)?;
 		}
 		let subject_buf = ctx.read_sandbox_memory(subject_ptr, subject_len)?;
@@ -1205,7 +1202,7 @@ define_env!(Env, <E: Ext>,
 			// allocator can handle.
 			ensure!(
 				delta_count
-					.saturating_mul(KEY_SIZE as u32) <= ctx.schedule.limits.max_memory_size(),
+					.saturating_mul(KEY_SIZE as u32) <= ctx.ext.schedule().limits.max_memory_size(),
 				Error::<E::T>::OutOfBounds,
 			);
 			let mut delta = vec![[0; KEY_SIZE]; delta_count as usize];
@@ -1253,7 +1250,7 @@ define_env!(Env, <E: Ext>,
 		};
 
 		// If there are more than `event_topics`, then trap.
-		if topics.len() > ctx.schedule.limits.event_topics as usize {
+		if topics.len() > ctx.ext.schedule().limits.event_topics as usize {
 			Err(Error::<E::T>::TooManyTopics)?;
 		}