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/benchmarking/code.rs

@@ -103,7 +103,7 @@ pub struct ImportedFunction {
 	pub return_type: Option<ValueType>,
 }
 
-/// A wasm module ready to be put on chain with `put_code`.
+/// A wasm module ready to be put on chain.
 #[derive(Clone)]
 pub struct WasmModule<T:Config> {
 	pub code: Vec<u8>,
@@ -245,16 +245,16 @@ where
 	}
 
 	/// Creates a wasm module of `target_bytes` size. Used to benchmark the performance of
-	/// `put_code` for different sizes of wasm modules. The generated module maximizes
+	/// `instantiate_with_code` for different sizes of wasm modules. The generated module maximizes
 	/// instrumentation runtime by nesting blocks as deeply as possible given the byte budget.
 	pub fn sized(target_bytes: u32) -> Self {
 		use parity_wasm::elements::Instruction::{If, I32Const, Return, End};
-		// Base size of a contract is 47 bytes and each expansion adds 6 bytes.
+		// Base size of a contract is 63 bytes and each expansion adds 6 bytes.
 		// We do one expansion less to account for the code section and function body
 		// size fields inside the binary wasm module representation which are leb128 encoded
 		// and therefore grow in size when the contract grows. We are not allowed to overshoot
-		// because of the maximum code size that is enforced by `put_code`.
-		let expansions = (target_bytes.saturating_sub(47) / 6).saturating_sub(1);
+		// because of the maximum code size that is enforced by `instantiate_with_code`.
+		let expansions = (target_bytes.saturating_sub(63) / 6).saturating_sub(1);
 		const EXPANSION: [Instruction; 4] = [
 			I32Const(0),
 			If(BlockType::NoResult),
@@ -263,6 +263,7 @@ where
 		];
 		ModuleDefinition {
 			call_body: Some(body::repeated(expansions, &EXPANSION)),
+			memory: Some(ImportedMemory::max::<T>()),
 			.. Default::default()
 		}
 		.into()

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

@@ -137,7 +137,7 @@ where
 		// same block number.
 		System::<T>::set_block_number(1u32.into());
 
-		Contracts::<T>::put_code_raw(module.code)?;
+		Contracts::<T>::store_code_raw(module.code)?;
 		Contracts::<T>::instantiate(
 			RawOrigin::Signed(caller.clone()).into(),
 			endowment,
@@ -198,7 +198,7 @@ where
 	/// Get the block number when this contract will be evicted. Returns an error when
 	/// the rent collection won't happen because the contract has to much endowment.
 	fn eviction_at(&self) -> Result<T::BlockNumber, &'static str> {
-		let projection = Rent::<T>::compute_projection(&self.account_id)
+		let projection = Rent::<T, PrefabWasmModule<T>>::compute_projection(&self.account_id)
 			.map_err(|_| "Invalid acc for rent")?;
 		match projection {
 			RentProjection::EvictionAt(at) => Ok(at),
@@ -250,7 +250,7 @@ where
 	/// Evict this contract.
 	fn evict(&mut self) -> Result<(), &'static str> {
 		self.set_block_num_for_eviction()?;
-		Rent::<T>::try_eviction(&self.contract.account_id, Zero::zero())?;
+		Rent::<T, PrefabWasmModule<T>>::try_eviction(&self.contract.account_id, Zero::zero())?;
 		self.contract.ensure_tombstone()
 	}
 }
@@ -314,24 +314,34 @@ benchmarks! {
 
 	// This constructs a contract that is maximal expensive to instrument.
 	// It creates a maximum number of metering blocks per byte.
-	// `n`: Size of the code in kilobytes.
-	put_code {
-		let n in 0 .. Contracts::<T>::current_schedule().limits.code_size / 1024;
+	// The size of the salt influences the runtime because is is hashed in order to
+	// determine the contract address.
+	// `c`: Size of the code in kilobytes.
+	// `s`: Size of the salt in kilobytes.
+	instantiate_with_code {
+		let c in 0 .. Contracts::<T>::current_schedule().limits.code_size / 1024;
+		let s in 0 .. code::max_pages::<T>() * 64;
+		let salt = vec![42u8; (s * 1024) as usize];
+		let endowment = caller_funding::<T>() / 3u32.into();
 		let caller = whitelisted_caller();
 		T::Currency::make_free_balance_be(&caller, caller_funding::<T>());
-		let module = WasmModule::<T>::sized(n * 1024);
-		let origin = RawOrigin::Signed(caller);
-	}: _(origin, module.code)
+		let WasmModule { code, hash, .. } = WasmModule::<T>::sized(c * 1024);
+		let origin = RawOrigin::Signed(caller.clone());
+		let addr = Contracts::<T>::contract_address(&caller, &hash, &salt);
+	}: _(origin, endowment, Weight::max_value(), code, vec![], salt)
+	verify {
+		// endowment was removed from the caller
+		assert_eq!(T::Currency::free_balance(&caller), caller_funding::<T>() - endowment);
+		// contract has the full endowment because no rent collection happended
+		assert_eq!(T::Currency::free_balance(&addr), endowment);
+		// instantiate should leave a alive contract
+		Contract::<T>::address_alive_info(&addr)?;
+	}
 
 	// Instantiate uses a dummy contract constructor to measure the overhead of the instantiate.
-	// The size of the input data influences the runtime because it is hashed in order to determine
-	// the contract address.
-	// `n`: Size of the data passed to constructor in kilobytes.
 	// `s`: Size of the salt in kilobytes.
 	instantiate {
-		let n in 0 .. code::max_pages::<T>() * 64;
 		let s in 0 .. code::max_pages::<T>() * 64;
-		let data = vec![42u8; (n * 1024) as usize];
 		let salt = vec![42u8; (s * 1024) as usize];
 		let endowment = caller_funding::<T>() / 3u32.into();
 		let caller = whitelisted_caller();
@@ -339,8 +349,8 @@ benchmarks! {
 		let WasmModule { code, hash, .. } = WasmModule::<T>::dummy_with_mem();
 		let origin = RawOrigin::Signed(caller.clone());
 		let addr = Contracts::<T>::contract_address(&caller, &hash, &salt);
-		Contracts::<T>::put_code_raw(code)?;
-	}: _(origin, endowment, Weight::max_value(), hash, data, salt)
+		Contracts::<T>::store_code_raw(code)?;
+	}: _(origin, endowment, Weight::max_value(), hash, vec![], salt)
 	verify {
 		// endowment was removed from the caller
 		assert_eq!(T::Currency::free_balance(&caller), caller_funding::<T>() - endowment);
@@ -1369,7 +1379,7 @@ benchmarks! {
 					])),
 					.. Default::default()
 				});
-				Contracts::<T>::put_code_raw(code.code)?;
+				Contracts::<T>::store_code_raw(code.code)?;
 				Ok(code.hash)
 			})
 			.collect::<Result<Vec<_>, &'static str>>()?;
@@ -1492,7 +1502,7 @@ benchmarks! {
 		let hash = callee_code.hash.clone();
 		let hash_bytes = callee_code.hash.encode();
 		let hash_len = hash_bytes.len();
-		Contracts::<T>::put_code_raw(callee_code.code)?;
+		Contracts::<T>::store_code_raw(callee_code.code)?;
 		let inputs = (0..API_BENCHMARK_BATCH_SIZE).map(|x| x.encode()).collect::<Vec<_>>();
 		let input_len = inputs.get(0).map(|x| x.len()).unwrap_or(0);
 		let input_bytes = inputs.iter().cloned().flatten().collect::<Vec<_>>();
@@ -2455,7 +2465,7 @@ mod tests {
 	create_test!(on_initialize_per_queue_item);
 
 	create_test!(update_schedule);
-	create_test!(put_code);
+	create_test!(instantiate_with_code);
 	create_test!(instantiate);
 	create_test!(call);
 	create_test!(claim_surcharge);

substrate/frame/contracts/src/gas.rs

@@ -18,8 +18,9 @@
 use crate::Config;
 use sp_std::marker::PhantomData;
 use sp_runtime::traits::Zero;
-use frame_support::dispatch::{
-	DispatchResultWithPostInfo, PostDispatchInfo, DispatchErrorWithPostInfo,
+use frame_support::{
+	dispatch::{DispatchResultWithPostInfo, PostDispatchInfo, DispatchErrorWithPostInfo},
+	weights::Weight,
 };
 use pallet_contracts_primitives::ExecError;
 
@@ -27,7 +28,7 @@ use pallet_contracts_primitives::ExecError;
 use std::{any::Any, fmt::Debug};
 
 // Gas is essentially the same as weight. It is a 1 to 1 correspondence.
-pub type Gas = frame_support::weights::Weight;
+pub type Gas = Weight;
 
 #[must_use]
 #[derive(Debug, PartialEq, Eq)]
@@ -201,12 +202,15 @@ impl<T: Config> GasMeter<T> {
 	}
 
 	/// Turn this GasMeter into a DispatchResult that contains the actually used gas.
-	pub fn into_dispatch_result<R, E>(self, result: Result<R, E>) -> DispatchResultWithPostInfo
+	pub fn into_dispatch_result<R, E>(
+		self, result: Result<R, E>,
+		base_weight: Weight,
+	) -> DispatchResultWithPostInfo
 	where
 		E: Into<ExecError>,
 	{
 		let post_info = PostDispatchInfo {
-			actual_weight: Some(self.gas_spent()),
+			actual_weight: Some(self.gas_spent().saturating_add(base_weight)),
 			pays_fee: Default::default(),
 		};
 

substrate/frame/contracts/src/lib.rs

@@ -59,10 +59,11 @@
 //!
 //! ### Dispatchable functions
 //!
-//! * `put_code` - Stores the given binary Wasm code into the chain's storage and returns its `code_hash`.
-//! * `instantiate` - Deploys a new contract from the given `code_hash`, optionally transferring some balance.
-//! This instantiates a new smart contract account and calls its contract deploy handler to
-//! initialize the contract.
+//! * `instantiate_with_code` - Deploys a new contract from the supplied wasm binary, optionally transferring
+//! some balance. This instantiates a new smart contract account and calls its contract deploy
+//! handler to initialize the contract.
+//! * `instantiate` - The same as `instantiate_with_code` but instead of uploading new code an
+//! existing `code_hash` is supplied.
 //! * `call` - Makes a call to an account, optionally transferring some balance.
 //!
 //! ## Usage
@@ -98,13 +99,12 @@ mod tests;
 
 pub use crate::{
 	gas::{Gas, GasMeter},
-	wasm::ReturnCode as RuntimeReturnCode,
+	wasm::{ReturnCode as RuntimeReturnCode, PrefabWasmModule},
 	weights::WeightInfo,
 	schedule::{Schedule, HostFnWeights, InstructionWeights, Limits},
 };
 use crate::{
-	exec::ExecutionContext,
-	wasm::{WasmLoader, WasmVm},
+	exec::{ExecutionContext, Executable},
 	rent::Rent,
 	storage::Storage,
 };
@@ -387,7 +387,8 @@ decl_error! {
 		/// The contract that was called is either no contract at all (a plain account)
 		/// or is a tombstone.
 		NotCallable,
-		/// The code supplied to `put_code` exceeds the limit specified in the current schedule.
+		/// The code supplied to `instantiate_with_code` exceeds the limit specified in the
+		/// current schedule.
 		CodeTooLarge,
 		/// No code could be found at the supplied code hash.
 		CodeNotFound,
@@ -431,6 +432,8 @@ decl_error! {
 		/// This can either happen when the accumulated storage in bytes is too large or
 		/// when number of storage items is too large.
 		StorageExhausted,
+		/// A contract with the same AccountId already exists.
+		DuplicateContract,
 	}
 }
 
@@ -528,23 +531,6 @@ decl_module! {
 			Ok(())
 		}
 
-		/// Stores the given binary Wasm code into the chain's storage and returns its `codehash`.
-		/// You can instantiate contracts only with stored code.
-		#[weight = T::WeightInfo::put_code(code.len() as u32 / 1024)]
-		pub fn put_code(
-			origin,
-			code: Vec<u8>
-		) -> DispatchResult {
-			ensure_signed(origin)?;
-			let schedule = <Module<T>>::current_schedule();
-			ensure!(code.len() as u32 <= schedule.limits.code_size, Error::<T>::CodeTooLarge);
-			let result = wasm::save_code::<T>(code, &schedule);
-			if let Ok(code_hash) = result {
-				Self::deposit_event(RawEvent::CodeStored(code_hash));
-			}
-			result.map(|_| ()).map_err(Into::into)
-		}
-
 		/// Makes a call to an account, optionally transferring some balance.
 		///
 		/// * If the account is a smart-contract account, the associated code will be
@@ -563,31 +549,73 @@ decl_module! {
 			let origin = ensure_signed(origin)?;
 			let dest = T::Lookup::lookup(dest)?;
 			let mut gas_meter = GasMeter::new(gas_limit);
-
 			let result = Self::execute_wasm(origin, &mut gas_meter, |ctx, gas_meter| {
 				ctx.call(dest, value, gas_meter, data)
 			});
-			gas_meter.into_dispatch_result(result)
+			gas_meter.into_dispatch_result(result, T::WeightInfo::call())
 		}
 
-		/// Instantiates a new contract from the `code_hash` generated by `put_code`,
-		/// optionally transferring some balance.
+		/// Instantiates a new contract from the supplied `code` optionally transferring
+		/// some balance.
 		///
-		/// The supplied `salt` is used for contract address deriviation. See `fn contract_address`.
+		/// This is the only function that can deploy new code to the chain.
+		///
+		/// # Parameters
+		///
+		/// * `endowment`: The balance to transfer from the `origin` to the newly created contract.
+		/// * `gas_limit`: The gas limit enforced when executing the constructor.
+		/// * `code`: The contract code to deploy in raw bytes.
+		/// * `data`: The input data to pass to the contract constructor.
+		/// * `salt`: Used for the address derivation. See [`Self::contract_address`].
 		///
 		/// Instantiation is executed as follows:
 		///
+		/// - The supplied `code` is instrumented, deployed, and a `code_hash` is created for that code.
+		/// - If the `code_hash` already exists on the chain the underlying `code` will be shared.
 		/// - The destination address is computed based on the sender, code_hash and the salt.
 		/// - The smart-contract account is created at the computed address.
-		/// - The `ctor_code` is executed in the context of the newly-created account. Buffer returned
-		///   after the execution is saved as the `code` of the account. That code will be invoked
-		///   upon any call received by this account.
-		/// - The contract is initialized.
+		/// - The `endowment` is transferred to the new account.
+		/// - The `deploy` function is executed in the context of the newly-created account.
 		#[weight =
-			T::WeightInfo::instantiate(
-				data.len() as u32 / 1024,
+			T::WeightInfo::instantiate_with_code(
+				code.len() as u32 / 1024,
 				salt.len() as u32 / 1024,
-			).saturating_add(*gas_limit)
+			)
+			.saturating_add(*gas_limit)
+		]
+		pub fn instantiate_with_code(
+			origin,
+			#[compact] endowment: BalanceOf<T>,
+			#[compact] gas_limit: Gas,
+			code: Vec<u8>,
+			data: Vec<u8>,
+			salt: Vec<u8>,
+		) -> DispatchResultWithPostInfo {
+			let origin = ensure_signed(origin)?;
+			let schedule = <Module<T>>::current_schedule();
+			let code_len = code.len() as u32;
+			ensure!(code_len <= schedule.limits.code_size, Error::<T>::CodeTooLarge);
+			let mut gas_meter = GasMeter::new(gas_limit);
+			let result = Self::execute_wasm(origin, &mut gas_meter, |ctx, gas_meter| {
+				let executable = PrefabWasmModule::from_code(code, &schedule)?;
+				let result = ctx.instantiate(endowment, gas_meter, executable, data, &salt)
+					.map(|(_address, output)| output)?;
+				Ok(result)
+			});
+			gas_meter.into_dispatch_result(
+				result,
+				T::WeightInfo::instantiate_with_code(code_len / 1024, salt.len() as u32 / 1024)
+			)
+		}
+
+		/// Instantiates a contract from a previously deployed wasm binary.
+		///
+		/// This function is identical to [`Self::instantiate_with_code`] but without the
+		/// code deployment step. Instead, the `code_hash` of an on-chain deployed wasm binary
+		/// must be supplied.
+		#[weight =
+			T::WeightInfo::instantiate(salt.len() as u32 / 1024)
+			.saturating_add(*gas_limit)
 		]
 		pub fn instantiate(
 			origin,
@@ -599,12 +627,16 @@ decl_module! {
 		) -> DispatchResultWithPostInfo {
 			let origin = ensure_signed(origin)?;
 			let mut gas_meter = GasMeter::new(gas_limit);
-
 			let result = Self::execute_wasm(origin, &mut gas_meter, |ctx, gas_meter| {
-				ctx.instantiate(endowment, gas_meter, &code_hash, data, &salt)
-					.map(|(_address, output)| output)
+				let executable = PrefabWasmModule::from_storage(code_hash, &ctx.config.schedule)?;
+				let result = ctx.instantiate(endowment, gas_meter, executable, data, &salt)
+					.map(|(_address, output)| output)?;
+				Ok(result)
 			});
-			gas_meter.into_dispatch_result(result)
+			gas_meter.into_dispatch_result(
+				result,
+				T::WeightInfo::instantiate(salt.len() as u32 / 1024)
+			)
 		}
 
 		/// Allows block producers to claim a small reward for evicting a contract. If a block
@@ -643,7 +675,9 @@ decl_module! {
 			};
 
 			// If poking the contract has lead to eviction of the contract, give out the rewards.
-			if let Some(rent_payed) = Rent::<T>::try_eviction(&dest, handicap)? {
+			if let Some(rent_payed) =
+				Rent::<T, PrefabWasmModule<T>>::try_eviction(&dest, handicap)?
+			{
 				T::Currency::deposit_into_existing(
 					&rewarded,
 					T::SurchargeReward::get().min(rent_payed),
@@ -698,15 +732,15 @@ where
 	}
 
 	pub fn rent_projection(address: T::AccountId) -> RentProjectionResult<T::BlockNumber> {
-		Rent::<T>::compute_projection(&address)
+		Rent::<T, PrefabWasmModule<T>>::compute_projection(&address)
 	}
 
-	/// Put code for benchmarks which does not check or instrument the code.
+	/// Store code for benchmarks which does not check nor instrument the code.
 	#[cfg(feature = "runtime-benchmarks")]
-	pub fn put_code_raw(code: Vec<u8>) -> DispatchResult {
+	pub fn store_code_raw(code: Vec<u8>) -> DispatchResult {
 		let schedule = <Module<T>>::current_schedule();
-		let result = wasm::save_code_raw::<T>(code, &schedule);
-		result.map(|_| ()).map_err(Into::into)
+		PrefabWasmModule::store_code_unchecked(code, &schedule)?;
+		Ok(())
 	}
 
 	/// Determine the address of a contract,
@@ -739,12 +773,13 @@ where
 	fn execute_wasm(
 		origin: T::AccountId,
 		gas_meter: &mut GasMeter<T>,
-		func: impl FnOnce(&mut ExecutionContext<T, WasmVm<T>, WasmLoader<T>>, &mut GasMeter<T>) -> ExecResult,
+		func: impl FnOnce(
+			&mut ExecutionContext<T, PrefabWasmModule<T>>,
+			&mut GasMeter<T>,
+		) -> ExecResult,
 	) -> ExecResult {
 		let cfg = ConfigCache::preload();
-		let vm = WasmVm::new(&cfg.schedule);
-		let loader = WasmLoader::new(&cfg.schedule);
-		let mut ctx = ExecutionContext::top_level(origin, &cfg, &vm, &loader);
+		let mut ctx = ExecutionContext::top_level(origin, &cfg);
 		func(&mut ctx, gas_meter)
 	}
 }
@@ -807,6 +842,12 @@ decl_event! {
 		/// - `data`: Data supplied by the contract. Metadata generated during contract
 		///           compilation is needed to decode it.
 		ContractEmitted(AccountId, Vec<u8>),
+
+		/// A code with the specified hash was removed.
+		/// \[code_hash\]
+		///
+		/// This happens when the last contract that uses this code hash was removed or evicted.
+		CodeRemoved(Hash),
 	}
 }
 
@@ -820,7 +861,7 @@ decl_storage! {
 		/// A mapping from an original code hash to the original code, untouched by instrumentation.
 		pub PristineCode: map hasher(identity) CodeHash<T> => Option<Vec<u8>>;
 		/// A mapping between an original code hash and instrumented wasm code, ready for execution.
-		pub CodeStorage: map hasher(identity) CodeHash<T> => Option<wasm::PrefabWasmModule>;
+		pub CodeStorage: map hasher(identity) CodeHash<T> => Option<PrefabWasmModule<T>>;
 		/// The subtrie counter.
 		pub AccountCounter: u64 = 0;
 		/// The code associated with a given account.

substrate/frame/contracts/src/rent.rs

@@ -20,7 +20,7 @@
 use crate::{
 	AliveContractInfo, BalanceOf, ContractInfo, ContractInfoOf, Module, RawEvent,
 	TombstoneContractInfo, Config, CodeHash, ConfigCache, Error,
-	storage::Storage,
+	storage::Storage, wasm::PrefabWasmModule, exec::Executable,
 };
 use sp_std::prelude::*;
 use sp_io::hashing::blake2_256;
@@ -86,12 +86,13 @@ enum Verdict<T: Config> {
 	Charge { amount: OutstandingAmount<T> },
 }
 
-pub struct Rent<T>(sp_std::marker::PhantomData<T>);
+pub struct Rent<T, E>(sp_std::marker::PhantomData<(T, E)>);
 
-impl<T> Rent<T>
+impl<T, E> Rent<T, E>
 where
 	T: Config,
-	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>
+	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>,
+	E: Executable<T>,
 {
 	/// Returns a fee charged per block from the contract.
 	///
@@ -99,10 +100,11 @@ where
 	/// then the fee can drop to zero.
 	fn compute_fee_per_block(
 		free_balance: &BalanceOf<T>,
-		contract: &AliveContractInfo<T>
+		contract: &AliveContractInfo<T>,
+		code_size_share: u32,
 	) -> BalanceOf<T> {
 		let uncovered_by_balance = T::DepositPerStorageByte::get()
-			.saturating_mul(contract.storage_size.into())
+			.saturating_mul(contract.storage_size.saturating_add(code_size_share).into())
 			.saturating_add(
 				T::DepositPerStorageItem::get()
 					.saturating_mul(contract.pair_count.into())
@@ -148,6 +150,7 @@ where
 		current_block_number: T::BlockNumber,
 		handicap: T::BlockNumber,
 		contract: &AliveContractInfo<T>,
+		code_size: u32,
 	) -> Verdict<T> {
 		// How much block has passed since the last deduction for the contract.
 		let blocks_passed = {
@@ -164,7 +167,7 @@ where
 		let free_balance = T::Currency::free_balance(account);
 
 		// An amount of funds to charge per block for storage taken up by the contract.
-		let fee_per_block = Self::compute_fee_per_block(&free_balance, contract);
+		let fee_per_block = Self::compute_fee_per_block(&free_balance, contract, code_size);
 		if fee_per_block.is_zero() {
 			// The rent deposit offset reduced the fee to 0. This means that the contract
 			// gets the rent for free.
@@ -228,19 +231,22 @@ where
 	/// Enacts the given verdict and returns the updated `ContractInfo`.
 	///
 	/// `alive_contract_info` should be from the same address as `account`.
+	///
+	/// # Note
+	///
+	/// if `evictable_code` is `None` an `Evict` verdict will not be enacted. This is for
+	/// when calling this function during a `call` where access to the soon to be evicted
+	/// contract should be denied but storage should be left unmodified.
 	fn enact_verdict(
 		account: &T::AccountId,
 		alive_contract_info: AliveContractInfo<T>,
 		current_block_number: T::BlockNumber,
 		verdict: Verdict<T>,
-		allow_eviction: bool,
-	) -> Result<Option<ContractInfo<T>>, DispatchError> {
-		match verdict {
-			Verdict::Exempt => return Ok(Some(ContractInfo::Alive(alive_contract_info))),
-			Verdict::Evict { amount: _ } if !allow_eviction => {
-				Ok(None)
-			}
-			Verdict::Evict { amount } => {
+		evictable_code: Option<PrefabWasmModule<T>>,
+	) -> Result<Option<AliveContractInfo<T>>, DispatchError> {
+		match (verdict, evictable_code) {
+			(Verdict::Exempt, _) => return Ok(Some(alive_contract_info)),
+			(Verdict::Evict { amount }, Some(code)) => {
 				// We need to remove the trie first because it is the only operation
 				// that can fail and this function is called without a storage
 				// transaction when called through `claim_surcharge`.
@@ -261,19 +267,23 @@ where
 				);
 				let tombstone_info = ContractInfo::Tombstone(tombstone);
 				<ContractInfoOf<T>>::insert(account, &tombstone_info);
+				code.drop_from_storage();
 				<Module<T>>::deposit_event(RawEvent::Evicted(account.clone()));
-				Ok(Some(tombstone_info))
+				Ok(None)
 			}
-			Verdict::Charge { amount } => {
-				let contract_info = ContractInfo::Alive(AliveContractInfo::<T> {
+			(Verdict::Evict { amount: _ }, None) => {
+				Ok(None)
+			}
+			(Verdict::Charge { amount }, _) => {
+				let contract = ContractInfo::Alive(AliveContractInfo::<T> {
 					rent_allowance: alive_contract_info.rent_allowance - amount.peek(),
 					deduct_block: current_block_number,
 					rent_payed: alive_contract_info.rent_payed.saturating_add(amount.peek()),
 					..alive_contract_info
 				});
-				<ContractInfoOf<T>>::insert(account, &contract_info);
+				<ContractInfoOf<T>>::insert(account, &contract);
 				amount.withdraw(account);
-				Ok(Some(contract_info))
+				Ok(Some(contract.get_alive().expect("We just constructed it as alive. qed")))
 			}
 		}
 	}
@@ -283,21 +293,20 @@ where
 	/// This functions does **not** evict the contract. It returns `None` in case the
 	/// contract is in need of eviction. [`try_eviction`] must
 	/// be called to perform the eviction.
-	pub fn charge(account: &T::AccountId) -> Result<Option<ContractInfo<T>>, DispatchError> {
-		let contract_info = <ContractInfoOf<T>>::get(account);
-		let alive_contract_info = match contract_info {
-			None | Some(ContractInfo::Tombstone(_)) => return Ok(contract_info),
-			Some(ContractInfo::Alive(contract)) => contract,
-		};
-
+	pub fn charge(
+		account: &T::AccountId,
+		contract: AliveContractInfo<T>,
+		code_size: u32,
+	) -> Result<Option<AliveContractInfo<T>>, DispatchError> {
 		let current_block_number = <frame_system::Module<T>>::block_number();
 		let verdict = Self::consider_case(
 			account,
 			current_block_number,
 			Zero::zero(),
-			&alive_contract_info,
+			&contract,
+			code_size,
 		);
-		Self::enact_verdict(account, alive_contract_info, current_block_number, verdict, false)
+		Self::enact_verdict(account, contract, current_block_number, verdict, None)
 	}
 
 	/// Process a report that a contract under the given address should be evicted.
@@ -322,12 +331,14 @@ where
 			None | Some(ContractInfo::Tombstone(_)) => return Ok(None),
 			Some(ContractInfo::Alive(contract)) => contract,
 		};
+		let module = PrefabWasmModule::<T>::from_storage_noinstr(contract.code_hash)?;
 		let current_block_number = <frame_system::Module<T>>::block_number();
 		let verdict = Self::consider_case(
 			account,
 			current_block_number,
 			handicap,
 			&contract,
+			module.occupied_storage(),
 		);
 
 		// Enact the verdict only if the contract gets removed.
@@ -339,7 +350,9 @@ where
 					.map(|a| a.peek())
 					.unwrap_or_else(|| <BalanceOf<T>>::zero())
 					.saturating_add(contract.rent_payed);
-				Self::enact_verdict(account, contract, current_block_number, verdict, true)?;
+				Self::enact_verdict(
+					account, contract, current_block_number, verdict, Some(module),
+				)?;
 				Ok(Some(rent_payed))
 			}
 			_ => Ok(None),
@@ -367,26 +380,33 @@ where
 			None | Some(ContractInfo::Tombstone(_)) => return Err(IsTombstone),
 			Some(ContractInfo::Alive(contract)) => contract,
 		};
+		let module = PrefabWasmModule::from_storage_noinstr(alive_contract_info.code_hash)
+			.map_err(|_| IsTombstone)?;
+		let code_size = module.occupied_storage();
 		let current_block_number = <frame_system::Module<T>>::block_number();
 		let verdict = Self::consider_case(
 			account,
 			current_block_number,
 			Zero::zero(),
 			&alive_contract_info,
+			code_size,
+		);
+		let new_contract_info = Self::enact_verdict(
+			account, alive_contract_info, current_block_number, verdict, Some(module),
 		);
-		let new_contract_info =
-			Self::enact_verdict(account, alive_contract_info, current_block_number, verdict, false);
 
 		// Check what happened after enaction of the verdict.
 		let alive_contract_info = match new_contract_info.map_err(|_| IsTombstone)? {
-			None | Some(ContractInfo::Tombstone(_)) => return Err(IsTombstone),
-			Some(ContractInfo::Alive(contract)) => contract,
+			None => return Err(IsTombstone),
+			Some(contract) => contract,
 		};
 
 		// Compute how much would the fee per block be with the *updated* balance.
 		let total_balance = T::Currency::total_balance(account);
 		let free_balance = T::Currency::free_balance(account);
-		let fee_per_block = Self::compute_fee_per_block(&free_balance, &alive_contract_info);
+		let fee_per_block = Self::compute_fee_per_block(
+			&free_balance, &alive_contract_info, code_size,
+		);
 		if fee_per_block.is_zero() {
 			return Ok(RentProjection::NoEviction);
 		}
@@ -418,6 +438,7 @@ where
 	/// Restores the destination account using the origin as prototype.
 	///
 	/// The restoration will be performed iff:
+	/// - the supplied code_hash does still exist on-chain
 	/// - origin exists and is alive,
 	/// - the origin's storage is not written in the current block
 	/// - the restored account has tombstone
@@ -455,6 +476,9 @@ where
 			origin_contract.last_write
 		};
 
+		// Fails if the code hash does not exist on chain
+		E::add_user(code_hash)?;
+
 		// We are allowed to eagerly modify storage even though the function can
 		// fail later due to tombstones not matching. This is because the restoration
 		// is always called from a contract and therefore in a storage transaction.
@@ -483,6 +507,7 @@ where
 		origin_contract.storage_size -= bytes_taken;
 
 		<ContractInfoOf<T>>::remove(&origin);
+		E::remove_user(origin_contract.code_hash);
 		<ContractInfoOf<T>>::insert(&dest, ContractInfo::Alive(AliveContractInfo::<T> {
 			trie_id: origin_contract.trie_id,
 			storage_size: origin_contract.storage_size,

substrate/frame/contracts/src/schedule.rs

@@ -106,7 +106,7 @@ pub struct Limits {
 	pub subject_len: u32,
 
 	/// The maximum length of a contract code in bytes. This limit applies to the uninstrumented
-	/// and pristine form of the code as supplied to `put_code`.
+	/// and pristine form of the code as supplied to `instantiate_with_code`.
 	pub code_size: u32,
 }
 

substrate/frame/contracts/src/storage.rs

@@ -164,29 +164,28 @@ where
 		account: &AccountIdOf<T>,
 		trie_id: TrieId,
 		ch: CodeHash<T>,
-	) -> Result<(), &'static str> {
-		<ContractInfoOf<T>>::mutate(account, |maybe_contract_info| {
-			if maybe_contract_info.is_some() {
-				return Err("Alive contract or tombstone already exists");
+	) -> DispatchResult {
+		<ContractInfoOf<T>>::try_mutate(account, |existing| {
+			if existing.is_some() {
+				return Err(Error::<T>::DuplicateContract.into());
 			}
 
-			*maybe_contract_info = Some(
-				AliveContractInfo::<T> {
-					code_hash: ch,
-					storage_size: 0,
-					trie_id,
-					deduct_block:
-						// We want to charge rent for the first block in advance. Therefore we
-						// treat the contract as if it was created in the last block and then
-						// charge rent for it during instantiation.
-						<frame_system::Module<T>>::block_number().saturating_sub(1u32.into()),
-					rent_allowance: <BalanceOf<T>>::max_value(),
-					rent_payed: <BalanceOf<T>>::zero(),
-					pair_count: 0,
-					last_write: None,
-				}
-				.into(),
-			);
+			let contract = AliveContractInfo::<T> {
+				code_hash: ch,
+				storage_size: 0,
+				trie_id,
+				deduct_block:
+					// We want to charge rent for the first block in advance. Therefore we
+					// treat the contract as if it was created in the last block and then
+					// charge rent for it during instantiation.
+					<frame_system::Module<T>>::block_number().saturating_sub(1u32.into()),
+				rent_allowance: <BalanceOf<T>>::max_value(),
+				rent_payed: <BalanceOf<T>>::zero(),
+				pair_count: 0,
+				last_write: None,
+			};
+
+			*existing = Some(contract.into());
 
 			Ok(())
 		})

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)?;
 		}