contracts: Consider contract size in weights (#8086)

* contracts: Consider contract size in weights

* Bump spec version

* Whitespace fix

Co-authored-by: Guillaume Thiolliere <gui.thiolliere@gmail.com>

* Correct pre-charged code weight even in the error case

* Use the instrumented code size in weight calculation

* Charge the cost of re-instrumentation from the gas meter

* Fix benchmark

* cargo run --release --features=runtime-benchmarks --manifest-path=bin/node/cli/Cargo.toml -- benchmark --chain=dev --steps=50 --repeat=20 --pallet=pallet_contracts --extrinsic=* --execution=wasm --wasm-execution=compiled --heap-pages=4096 --output=./frame/contracts/src/weights.rs --template=./.maintain/frame-weight-template.hbs

* Better documentation of return types

Co-authored-by: Guillaume Thiolliere <gui.thiolliere@gmail.com>
Co-authored-by: Parity Benchmarking Bot <admin@parity.io>

substrate/bin/node/runtime/src/lib.rs

@@ -112,8 +112,8 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
 	// and set impl_version to 0. If only runtime
 	// implementation changes and behavior does not, then leave spec_version as
 	// is and increment impl_version.
-	spec_version: 264,
-	impl_version: 1,
+	spec_version: 265,
+	impl_version: 0,
 	apis: RUNTIME_API_VERSIONS,
 	transaction_version: 2,
 };
@@ -735,6 +735,7 @@ parameter_types! {
 			<Runtime as pallet_contracts::Config>::WeightInfo::on_initialize_per_queue_item(1) -
 			<Runtime as pallet_contracts::Config>::WeightInfo::on_initialize_per_queue_item(0)
 		)) / 5) as u32;
+	pub MaxCodeSize: u32 = 128 * 1024;
 }
 
 impl pallet_contracts::Config for Runtime {
@@ -757,6 +758,7 @@ impl pallet_contracts::Config for Runtime {
 	type ChainExtension = ();
 	type DeletionQueueDepth = DeletionQueueDepth;
 	type DeletionWeightLimit = DeletionWeightLimit;
+	type MaxCodeSize = MaxCodeSize;
 }
 
 impl pallet_sudo::Config for Runtime {

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

@@ -27,12 +27,14 @@
 use crate::Config;
 use crate::Module as Contracts;
 
-use parity_wasm::elements::{Instruction, Instructions, FuncBody, ValueType, BlockType};
+use parity_wasm::elements::{
+	Instruction, Instructions, FuncBody, ValueType, BlockType, Section, CustomSection,
+};
 use pwasm_utils::stack_height::inject_limiter;
 use sp_core::crypto::UncheckedFrom;
 use sp_runtime::traits::Hash;
 use sp_sandbox::{EnvironmentDefinitionBuilder, Memory};
-use sp_std::{prelude::*, convert::TryFrom};
+use sp_std::{prelude::*, convert::TryFrom, borrow::ToOwned};
 
 /// Pass to `create_code` in order to create a compiled `WasmModule`.
 ///
@@ -66,6 +68,10 @@ pub struct ModuleDefinition {
 	pub inject_stack_metering: bool,
 	/// Create a table containing function pointers.
 	pub table: Option<TableSegment>,
+	/// Create a section named "dummy" of the specified size. This is useful in order to
+	/// benchmark the overhead of loading and storing codes of specified sizes. The dummy
+	/// section only contributes to the size of the contract but does not affect execution.
+	pub dummy_section: u32,
 }
 
 pub struct TableSegment {
@@ -204,6 +210,15 @@ where
 				.build();
 		}
 
+		// Add the dummy section
+		if def.dummy_section > 0 {
+			contract = contract.with_section(
+				Section::Custom(
+					CustomSection::new("dummy".to_owned(), vec![42; def.dummy_section as usize])
+				)
+			);
+		}
+
 		let mut code = contract.build();
 
 		// Inject stack height metering
@@ -235,10 +250,11 @@ where
 		ModuleDefinition::default().into()
 	}
 
-	/// Same as `dummy` but with maximum sized linear memory.
-	pub fn dummy_with_mem() -> Self {
+	/// Same as `dummy` but with maximum sized linear memory and a dummy section of specified size.
+	pub fn dummy_with_bytes(dummy_bytes: u32) -> Self {
 		ModuleDefinition {
 			memory: Some(ImportedMemory::max::<T>()),
+			dummy_section: dummy_bytes,
 			.. Default::default()
 		}
 		.into()

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

@@ -304,6 +304,19 @@ benchmarks! {
 		Storage::<T>::process_deletion_queue_batch(Weight::max_value())
 	}
 
+	// This benchmarks the additional weight that is charged when a contract is executed the
+	// first time after a new schedule was deployed: For every new schedule a contract needs
+	// to re-run the instrumentation once.
+	instrument {
+		let c in 0 .. T::MaxCodeSize::get() / 1024;
+		let WasmModule { code, hash, .. } = WasmModule::<T>::sized(c * 1024);
+		Contracts::<T>::store_code_raw(code)?;
+		let mut module = PrefabWasmModule::from_storage_noinstr(hash)?;
+		let schedule = Contracts::<T>::current_schedule();
+	}: {
+		Contracts::<T>::reinstrument_module(&mut module, &schedule)?;
+	}
+
 	// This extrinsic is pretty much constant as it is only a simple setter.
 	update_schedule {
 		let schedule = Schedule {
@@ -318,8 +331,13 @@ benchmarks! {
 	// determine the contract address.
 	// `c`: Size of the code in kilobytes.
 	// `s`: Size of the salt in kilobytes.
+	//
+	// # Note
+	//
+	// We cannot let `c` grow to the maximum code size because the code is not allowed
+	// to be larger than the maximum size **after instrumentation**.
 	instantiate_with_code {
-		let c in 0 .. Contracts::<T>::current_schedule().limits.code_size / 1024;
+		let c in 0 .. Perbill::from_percent(50).mul_ceil(T::MaxCodeSize::get() / 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();
@@ -339,14 +357,16 @@ benchmarks! {
 	}
 
 	// Instantiate uses a dummy contract constructor to measure the overhead of the instantiate.
+	// `c`: Size of the code in kilobytes.
 	// `s`: Size of the salt in kilobytes.
 	instantiate {
+		let c in 0 .. T::MaxCodeSize::get() / 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 WasmModule { code, hash, .. } = WasmModule::<T>::dummy_with_mem();
+		let WasmModule { code, hash, .. } = WasmModule::<T>::dummy_with_bytes(c * 1024);
 		let origin = RawOrigin::Signed(caller.clone());
 		let addr = Contracts::<T>::contract_address(&caller, &hash, &salt);
 		Contracts::<T>::store_code_raw(code)?;
@@ -365,10 +385,12 @@ benchmarks! {
 	// won't call `seal_input` in its constructor to copy the data to contract memory.
 	// The dummy contract used here does not do this. The costs for the data copy is billed as
 	// part of `seal_input`.
+	// `c`: Size of the code in kilobytes.
 	call {
+		let c in 0 .. T::MaxCodeSize::get() / 1024;
 		let data = vec![42u8; 1024];
 		let instance = Contract::<T>::with_caller(
-			whitelisted_caller(), WasmModule::dummy_with_mem(), vec![], Endow::CollectRent
+			whitelisted_caller(), WasmModule::dummy_with_bytes(c * 1024), vec![], Endow::CollectRent
 		)?;
 		let value = T::Currency::minimum_balance() * 100u32.into();
 		let origin = RawOrigin::Signed(instance.caller.clone());
@@ -396,9 +418,11 @@ benchmarks! {
 	// will be distributed over multiple blocks using a scheduler. Otherwise there is
 	// no incentive to remove large contracts when the removal is more expensive than
 	// the reward for removing them.
+	// `c`: Size of the code of the contract that should be evicted.
 	claim_surcharge {
+		let c in 0 .. T::MaxCodeSize::get() / 1024;
 		let instance = Contract::<T>::with_caller(
-			whitelisted_caller(), WasmModule::dummy(), vec![], Endow::CollectRent
+			whitelisted_caller(), WasmModule::dummy_with_bytes(c * 1024), vec![], Endow::CollectRent
 		)?;
 		let origin = RawOrigin::Signed(instance.caller.clone());
 		let account_id = instance.account_id.clone();
@@ -694,6 +718,42 @@ benchmarks! {
 		}
 	}
 
+	seal_terminate_per_code_kb {
+		let c in 0 .. T::MaxCodeSize::get() / 1024;
+		let beneficiary = account::<T::AccountId>("beneficiary", 0, 0);
+		let beneficiary_bytes = beneficiary.encode();
+		let beneficiary_len = beneficiary_bytes.len();
+		let code = WasmModule::<T>::from(ModuleDefinition {
+			memory: Some(ImportedMemory::max::<T>()),
+			imported_functions: vec![ImportedFunction {
+				name: "seal_terminate",
+				params: vec![ValueType::I32, ValueType::I32],
+				return_type: None,
+			}],
+			data_segments: vec![
+				DataSegment {
+					offset: 0,
+					value: beneficiary_bytes,
+				},
+			],
+			call_body: Some(body::repeated(1, &[
+				Instruction::I32Const(0), // beneficiary_ptr
+				Instruction::I32Const(beneficiary_len as i32), // beneficiary_len
+				Instruction::Call(0),
+			])),
+			dummy_section: c * 1024,
+			.. Default::default()
+		});
+		let instance = Contract::<T>::new(code, vec![], Endow::Max)?;
+		let origin = RawOrigin::Signed(instance.caller.clone());
+		assert_eq!(T::Currency::total_balance(&beneficiary), 0u32.into());
+		assert_eq!(T::Currency::total_balance(&instance.account_id), Endow::max::<T>());
+	}: call(origin, instance.addr, 0u32.into(), Weight::max_value(), vec![])
+	verify {
+		assert_eq!(T::Currency::total_balance(&instance.account_id), 0u32.into());
+		assert_eq!(T::Currency::total_balance(&beneficiary), Endow::max::<T>());
+	}
+
 	seal_restore_to {
 		let r in 0 .. 1;
 
@@ -772,9 +832,16 @@ benchmarks! {
 		}
 	}
 
-	seal_restore_to_per_delta {
+	// `c`: Code size of caller contract
+	// `t`: Code size of tombstone contract
+	// `d`: Number of supplied delta keys
+	seal_restore_to_per_code_kb_delta {
+		let c in 0 .. T::MaxCodeSize::get() / 1024;
+		let t in 0 .. T::MaxCodeSize::get() / 1024;
 		let d in 0 .. API_BENCHMARK_BATCHES;
-		let mut tombstone = ContractWithStorage::<T>::new(0, 0)?;
+		let mut tombstone = ContractWithStorage::<T>::with_code(
+			WasmModule::<T>::dummy_with_bytes(t * 1024), 0, 0
+		)?;
 		tombstone.evict()?;
 		let delta = create_storage::<T>(d * API_BENCHMARK_BATCH_SIZE, T::MaxValueSize::get())?;
 
@@ -837,6 +904,7 @@ benchmarks! {
 				Instruction::Call(0),
 				Instruction::End,
 			])),
+			dummy_section: c * 1024,
 			.. Default::default()
 		});
 
@@ -1225,7 +1293,7 @@ benchmarks! {
 	// We call unique accounts.
 	seal_call {
 		let r in 0 .. API_BENCHMARK_BATCHES;
-		let dummy_code = WasmModule::<T>::dummy_with_mem();
+		let dummy_code = WasmModule::<T>::dummy_with_bytes(0);
 		let callees = (0..r * API_BENCHMARK_BATCH_SIZE)
 			.map(|i| Contract::with_index(i + 1, dummy_code.clone(), vec![], Endow::Max))
 			.collect::<Result<Vec<_>, _>>()?;
@@ -1280,7 +1348,8 @@ benchmarks! {
 		let origin = RawOrigin::Signed(instance.caller.clone());
 	}: call(origin, instance.addr, 0u32.into(), Weight::max_value(), vec![])
 
-	seal_call_per_transfer_input_output_kb {
+	seal_call_per_code_transfer_input_output_kb {
+		let c in 0 .. T::MaxCodeSize::get() / 1024;
 		let t in 0 .. 1;
 		let i in 0 .. code::max_pages::<T>() * 64;
 		let o in 0 .. (code::max_pages::<T>() - 1) * 64;
@@ -1302,6 +1371,7 @@ benchmarks! {
 				Instruction::Call(0),
 				Instruction::End,
 			])),
+			dummy_section: c * 1024,
 			.. Default::default()
 		});
 		let callees = (0..API_BENCHMARK_BATCH_SIZE)
@@ -1475,7 +1545,8 @@ benchmarks! {
 		}
 	}
 
-	seal_instantiate_per_input_output_salt_kb {
+	seal_instantiate_per_code_input_output_salt_kb {
+		let c in 0 .. T::MaxCodeSize::get() / 1024;
 		let i in 0 .. (code::max_pages::<T>() - 1) * 64;
 		let o in 0 .. (code::max_pages::<T>() - 1) * 64;
 		let s in 0 .. (code::max_pages::<T>() - 1) * 64;
@@ -1497,6 +1568,7 @@ benchmarks! {
 				Instruction::Call(0),
 				Instruction::End,
 			])),
+			dummy_section: c * 1024,
 			.. Default::default()
 		});
 		let hash = callee_code.hash.clone();
@@ -2440,8 +2512,6 @@ benchmarks! {
 	}: {}
 }
 
-
-
 impl_benchmark_test_suite!(
 	Contracts,
 	crate::tests::ExtBuilder::default().build(),

substrate/frame/contracts/src/exec.rs

@@ -72,8 +72,13 @@ pub trait Ext {
 
 	/// Instantiate a contract from the given code.
 	///
+	/// Returns the original code size of the called contract.
 	/// The newly created account will be associated with `code`. `value` specifies the amount of value
 	/// transferred from this to the newly created account (also known as endowment).
+	///
+	/// # Return Value
+	///
+	/// Result<(AccountId, ExecReturnValue, CodeSize), (ExecError, CodeSize)>
 	fn instantiate(
 		&mut self,
 		code: CodeHash<Self::T>,
@@ -81,7 +86,7 @@ pub trait Ext {
 		gas_meter: &mut GasMeter<Self::T>,
 		input_data: Vec<u8>,
 		salt: &[u8],
-	) -> Result<(AccountIdOf<Self::T>, ExecReturnValue), ExecError>;
+	) -> Result<(AccountIdOf<Self::T>, ExecReturnValue, u32), (ExecError, u32)>;
 
 	/// Transfer some amount of funds into the specified account.
 	fn transfer(
@@ -92,24 +97,35 @@ pub trait Ext {
 
 	/// Transfer all funds to `beneficiary` and delete the contract.
 	///
+	/// Returns the original code size of the terminated contract.
 	/// Since this function removes the self contract eagerly, if succeeded, no further actions should
 	/// be performed on this `Ext` instance.
 	///
 	/// This function will fail if the same contract is present on the contract
 	/// call stack.
+	///
+	/// # Return Value
+	///
+	/// Result<CodeSize, (DispatchError, CodeSize)>
 	fn terminate(
 		&mut self,
 		beneficiary: &AccountIdOf<Self::T>,
-	) -> DispatchResult;
+	) -> Result<u32, (DispatchError, u32)>;
 
 	/// Call (possibly transferring some amount of funds) into the specified account.
+	///
+	/// Returns the original code size of the called contract.
+	///
+	/// # Return Value
+	///
+	/// Result<(ExecReturnValue, CodeSize), (ExecError, CodeSize)>
 	fn call(
 		&mut self,
 		to: &AccountIdOf<Self::T>,
 		value: BalanceOf<Self::T>,
 		gas_meter: &mut GasMeter<Self::T>,
 		input_data: Vec<u8>,
-	) -> ExecResult;
+	) -> Result<(ExecReturnValue, u32), (ExecError, u32)>;
 
 	/// Restores the given destination contract sacrificing the current one.
 	///
@@ -118,13 +134,17 @@ pub trait Ext {
 	///
 	/// This function will fail if the same contract is present
 	/// on the contract call stack.
+	///
+	/// # Return Value
+	///
+	/// Result<(CallerCodeSize, DestCodeSize), (DispatchError, CallerCodeSize, DestCodesize)>
 	fn restore_to(
 		&mut self,
 		dest: AccountIdOf<Self::T>,
 		code_hash: CodeHash<Self::T>,
 		rent_allowance: BalanceOf<Self::T>,
 		delta: Vec<StorageKey>,
-	) -> DispatchResult;
+	) -> Result<(u32, u32), (DispatchError, u32, u32)>;
 
 	/// Returns a reference to the account id of the caller.
 	fn caller(&self) -> &AccountIdOf<Self::T>;
@@ -190,7 +210,11 @@ pub enum ExportedFunction {
 /// order to be able to mock the wasm logic for testing.
 pub trait Executable<T: Config>: Sized {
 	/// Load the executable from storage.
-	fn from_storage(code_hash: CodeHash<T>, schedule: &Schedule<T>) -> Result<Self, DispatchError>;
+	fn from_storage(
+		code_hash: CodeHash<T>,
+		schedule: &Schedule<T>,
+		gas_meter: &mut GasMeter<T>,
+	) -> Result<Self, DispatchError>;
 
 	/// Load the module from storage without re-instrumenting it.
 	///
@@ -203,10 +227,14 @@ pub trait Executable<T: Config>: Sized {
 	fn drop_from_storage(self);
 
 	/// Increment the refcount by one. Fails if the code does not exist on-chain.
-	fn add_user(code_hash: CodeHash<T>) -> DispatchResult;
+	///
+	/// Returns the size of the original code.
+	fn add_user(code_hash: CodeHash<T>) -> Result<u32, DispatchError>;
 
 	/// Decrement the refcount by one and remove the code when it drops to zero.
-	fn remove_user(code_hash: CodeHash<T>);
+	///
+	/// Returns the size of the original code.
+	fn remove_user(code_hash: CodeHash<T>) -> u32;
 
 	/// Execute the specified exported function and return the result.
 	///
@@ -238,6 +266,9 @@ pub trait Executable<T: Config>: Sized {
 	/// without refetching this from storage the result can be inaccurate as it might be
 	/// working with a stale value. Usually this inaccuracy is tolerable.
 	fn occupied_storage(&self) -> u32;
+
+	/// Size of the instrumented code in bytes.
+	fn code_len(&self) -> u32;
 }
 
 pub struct ExecutionContext<'a, T: Config + 'a, E> {
@@ -290,35 +321,42 @@ where
 	}
 
 	/// Make a call to the specified address, optionally transferring some funds.
+	///
+	/// # Return Value
+	///
+	/// Result<(ExecReturnValue, CodeSize), (ExecError, CodeSize)>
 	pub fn call(
 		&mut self,
 		dest: T::AccountId,
 		value: BalanceOf<T>,
 		gas_meter: &mut GasMeter<T>,
 		input_data: Vec<u8>,
-	) -> ExecResult {
+	) -> Result<(ExecReturnValue, u32), (ExecError, u32)> {
 		if self.depth == T::MaxDepth::get() as usize {
-			Err(Error::<T>::MaxCallDepthReached)?
+			return Err((Error::<T>::MaxCallDepthReached.into(), 0));
 		}
 
 		let contract = <ContractInfoOf<T>>::get(&dest)
 			.and_then(|contract| contract.get_alive())
-			.ok_or(Error::<T>::NotCallable)?;
+			.ok_or((Error::<T>::NotCallable.into(), 0))?;
 
-		let executable = E::from_storage(contract.code_hash, &self.schedule)?;
+		let executable = E::from_storage(contract.code_hash, &self.schedule, gas_meter)
+			.map_err(|e| (e.into(), 0))?;
+		let code_len = executable.code_len();
 
 		// This charges the rent and denies access to a contract that is in need of
 		// eviction by returning `None`. We cannot evict eagerly here because those
 		// changes would be rolled back in case this contract is called by another
 		// contract.
 		// See: https://github.com/paritytech/substrate/issues/6439#issuecomment-648754324
-		let contract = Rent::<T, E>::charge(&dest, contract, executable.occupied_storage())?
-			.ok_or(Error::<T>::NotCallable)?;
+		let contract = Rent::<T, E>::charge(&dest, contract, executable.occupied_storage())
+			.map_err(|e| (e.into(), code_len))?
+			.ok_or((Error::<T>::NotCallable.into(), code_len))?;
 
 		let transactor_kind = self.transactor_kind();
 		let caller = self.self_account.clone();
 
-		self.with_nested_context(dest.clone(), contract.trie_id.clone(), |nested| {
+		let result = self.with_nested_context(dest.clone(), contract.trie_id.clone(), |nested| {
 			if value > BalanceOf::<T>::zero() {
 				transfer::<T>(
 					TransferCause::Call,
@@ -336,7 +374,8 @@ where
 				gas_meter,
 			).map_err(|e| ExecError { error: e.error, origin: ErrorOrigin::Callee })?;
 			Ok(output)
-		})
+		}).map_err(|e| (e, code_len))?;
+		Ok((result, code_len))
 	}
 
 	pub fn instantiate(
@@ -581,10 +620,13 @@ where
 		gas_meter: &mut GasMeter<T>,
 		input_data: Vec<u8>,
 		salt: &[u8],
-	) -> Result<(AccountIdOf<T>, ExecReturnValue), ExecError> {
-		let executable = E::from_storage(code_hash, &self.ctx.schedule)?;
-		let result = self.ctx.instantiate(endowment, gas_meter, executable, input_data, salt)?;
-		Ok(result)
+	) -> Result<(AccountIdOf<T>, ExecReturnValue, u32), (ExecError, u32)> {
+		let executable = E::from_storage(code_hash, &self.ctx.schedule, gas_meter)
+			.map_err(|e| (e.into(), 0))?;
+		let code_len = executable.code_len();
+		self.ctx.instantiate(endowment, gas_meter, executable, input_data, salt)
+			.map(|r| (r.0, r.1, code_len))
+			.map_err(|e| (e, code_len))
 	}
 
 	fn transfer(
@@ -604,12 +646,12 @@ where
 	fn terminate(
 		&mut self,
 		beneficiary: &AccountIdOf<Self::T>,
-	) -> DispatchResult {
+	) -> Result<u32, (DispatchError, u32)> {
 		let self_id = self.ctx.self_account.clone();
 		let value = T::Currency::free_balance(&self_id);
 		if let Some(caller_ctx) = self.ctx.caller {
 			if caller_ctx.is_live(&self_id) {
-				return Err(Error::<T>::ReentranceDenied.into());
+				return Err((Error::<T>::ReentranceDenied.into(), 0));
 			}
 		}
 		transfer::<T>(
@@ -618,12 +660,12 @@ where
 			&self_id,
 			beneficiary,
 			value,
-		)?;
+		).map_err(|e| (e, 0))?;
 		if let Some(ContractInfo::Alive(info)) = ContractInfoOf::<T>::take(&self_id) {
-			Storage::<T>::queue_trie_for_deletion(&info)?;
-			E::remove_user(info.code_hash);
+			Storage::<T>::queue_trie_for_deletion(&info).map_err(|e| (e, 0))?;
+			let code_len = E::remove_user(info.code_hash);
 			Contracts::<T>::deposit_event(RawEvent::Terminated(self_id, beneficiary.clone()));
-			Ok(())
+			Ok(code_len)
 		} else {
 			panic!(
 				"this function is only invoked by in the context of a contract;\
@@ -639,7 +681,7 @@ where
 		value: BalanceOf<T>,
 		gas_meter: &mut GasMeter<T>,
 		input_data: Vec<u8>,
-	) -> ExecResult {
+	) -> Result<(ExecReturnValue, u32), (ExecError, u32)> {
 		self.ctx.call(to.clone(), value, gas_meter, input_data)
 	}
 
@@ -649,10 +691,10 @@ where
 		code_hash: CodeHash<Self::T>,
 		rent_allowance: BalanceOf<Self::T>,
 		delta: Vec<StorageKey>,
-	) -> DispatchResult {
+	) -> Result<(u32, u32), (DispatchError, u32, u32)> {
 		if let Some(caller_ctx) = self.ctx.caller {
 			if caller_ctx.is_live(&self.ctx.self_account) {
-				return Err(Error::<T>::ReentranceDenied.into());
+				return Err((Error::<T>::ReentranceDenied.into(), 0, 0));
 			}
 		}
 
@@ -828,7 +870,8 @@ mod tests {
 	impl Executable<Test> for MockExecutable {
 		fn from_storage(
 			code_hash: CodeHash<Test>,
-			_schedule: &Schedule<Test>
+			_schedule: &Schedule<Test>,
+			_gas_meter: &mut GasMeter<Test>,
 		) -> Result<Self, DispatchError> {
 			Self::from_storage_noinstr(code_hash)
 		}
@@ -845,11 +888,11 @@ mod tests {
 
 		fn drop_from_storage(self) {}
 
-		fn add_user(_code_hash: CodeHash<Test>) -> DispatchResult {
-			Ok(())
+		fn add_user(_code_hash: CodeHash<Test>) -> Result<u32, DispatchError> {
+			Ok(0)
 		}
 
-		fn remove_user(_code_hash: CodeHash<Test>) {}
+		fn remove_user(_code_hash: CodeHash<Test>) -> u32 { 0 }
 
 		fn execute<E: Ext<T = Test>>(
 			self,
@@ -872,6 +915,10 @@ mod tests {
 		fn occupied_storage(&self) -> u32 {
 			0
 		}
+
+		fn code_len(&self) -> u32 {
+			0
+		}
 	}
 
 	fn exec_success() -> ExecResult {
@@ -954,7 +1001,7 @@ mod tests {
 				vec![],
 			).unwrap();
 
-			assert!(!output.is_success());
+			assert!(!output.0.is_success());
 			assert_eq!(get_balance(&origin), 100);
 
 			// the rent is still charged
@@ -1012,8 +1059,8 @@ mod tests {
 			);
 
 			let output = result.unwrap();
-			assert!(output.is_success());
-			assert_eq!(output.data, vec![1, 2, 3, 4]);
+			assert!(output.0.is_success());
+			assert_eq!(output.0.data, vec![1, 2, 3, 4]);
 		});
 	}
 
@@ -1040,8 +1087,8 @@ mod tests {
 			);
 
 			let output = result.unwrap();
-			assert!(!output.is_success());
-			assert_eq!(output.data, vec![1, 2, 3, 4]);
+			assert!(!output.0.is_success());
+			assert_eq!(output.0.data, vec![1, 2, 3, 4]);
 		});
 	}
 
@@ -1080,13 +1127,17 @@ mod tests {
 			let schedule = Contracts::current_schedule();
 			let subsistence = Contracts::<Test>::subsistence_threshold();
 			let mut ctx = MockContext::top_level(ALICE, &schedule);
+			let mut gas_meter = GasMeter::<Test>::new(GAS_LIMIT);
+			let executable = MockExecutable::from_storage(
+				input_data_ch, &schedule, &mut gas_meter
+			).unwrap();
 
 			set_balance(&ALICE, subsistence * 10);
 
 			let result = ctx.instantiate(
 				subsistence * 3,
-				&mut GasMeter::<Test>::new(GAS_LIMIT),
-				MockExecutable::from_storage(input_data_ch, &schedule).unwrap(),
+				&mut gas_meter,
+				executable,
 				vec![1, 2, 3, 4],
 				&[],
 			);
@@ -1113,7 +1164,7 @@ mod tests {
 					// Verify that we've got proper error and set `reached_bottom`.
 					assert_eq!(
 						r,
-						Err(Error::<Test>::MaxCallDepthReached.into())
+						Err((Error::<Test>::MaxCallDepthReached.into(), 0))
 					);
 					*reached_bottom = true;
 				} else {
@@ -1235,12 +1286,16 @@ mod tests {
 		ExtBuilder::default().existential_deposit(15).build().execute_with(|| {
 			let schedule = Contracts::current_schedule();
 			let mut ctx = MockContext::top_level(ALICE, &schedule);
+			let mut gas_meter = GasMeter::<Test>::new(GAS_LIMIT);
+			let executable = MockExecutable::from_storage(
+				dummy_ch, &schedule, &mut gas_meter
+			).unwrap();
 
 			assert_matches!(
 				ctx.instantiate(
 					0, // <- zero endowment
-					&mut GasMeter::<Test>::new(GAS_LIMIT),
-					MockExecutable::from_storage(dummy_ch, &schedule).unwrap(),
+					&mut gas_meter,
+					executable,
 					vec![],
 					&[],
 				),
@@ -1258,13 +1313,17 @@ mod tests {
 		ExtBuilder::default().existential_deposit(15).build().execute_with(|| {
 			let schedule = Contracts::current_schedule();
 			let mut ctx = MockContext::top_level(ALICE, &schedule);
+			let mut gas_meter = GasMeter::<Test>::new(GAS_LIMIT);
+			let executable = MockExecutable::from_storage(
+				dummy_ch, &schedule, &mut gas_meter
+			).unwrap();
 			set_balance(&ALICE, 1000);
 
 			let instantiated_contract_address = assert_matches!(
 				ctx.instantiate(
 					100,
-					&mut GasMeter::<Test>::new(GAS_LIMIT),
-					MockExecutable::from_storage(dummy_ch, &schedule).unwrap(),
+					&mut gas_meter,
+					executable,
 					vec![],
 					&[],
 				),
@@ -1289,13 +1348,17 @@ mod tests {
 		ExtBuilder::default().existential_deposit(15).build().execute_with(|| {
 			let schedule = Contracts::current_schedule();
 			let mut ctx = MockContext::top_level(ALICE, &schedule);
+			let mut gas_meter = GasMeter::<Test>::new(GAS_LIMIT);
+			let executable = MockExecutable::from_storage(
+				dummy_ch, &schedule, &mut gas_meter
+			).unwrap();
 			set_balance(&ALICE, 1000);
 
 			let instantiated_contract_address = assert_matches!(
 				ctx.instantiate(
 					100,
-					&mut GasMeter::<Test>::new(GAS_LIMIT),
-					MockExecutable::from_storage(dummy_ch, &schedule).unwrap(),
+					&mut gas_meter,
+					executable,
 					vec![],
 					&[],
 				),
@@ -1317,7 +1380,7 @@ mod tests {
 			let instantiated_contract_address = Rc::clone(&instantiated_contract_address);
 			move |ctx| {
 				// Instantiate a contract and save it's address in `instantiated_contract_address`.
-				let (address, output) = ctx.ext.instantiate(
+				let (address, output, _) = ctx.ext.instantiate(
 					dummy_ch,
 					Contracts::<Test>::subsistence_threshold() * 3,
 					ctx.gas_meter,
@@ -1369,10 +1432,10 @@ mod tests {
 						vec![],
 						&[],
 					),
-					Err(ExecError {
+					Err((ExecError {
 						error: DispatchError::Other("It's a trap!"),
 						origin: ErrorOrigin::Callee,
-					})
+					}, 0))
 				);
 
 				exec_success()
@@ -1410,13 +1473,17 @@ mod tests {
 			.execute_with(|| {
 				let schedule = Contracts::current_schedule();
 				let mut ctx = MockContext::top_level(ALICE, &schedule);
+				let mut gas_meter = GasMeter::<Test>::new(GAS_LIMIT);
+				let executable = MockExecutable::from_storage(
+					terminate_ch, &schedule, &mut gas_meter
+				).unwrap();
 				set_balance(&ALICE, 1000);
 
 				assert_eq!(
 					ctx.instantiate(
 						100,
-						&mut GasMeter::<Test>::new(GAS_LIMIT),
-						MockExecutable::from_storage(terminate_ch, &schedule).unwrap(),
+						&mut gas_meter,
+						executable,
 						vec![],
 						&[],
 					),
@@ -1445,12 +1512,16 @@ mod tests {
 			let subsistence = Contracts::<Test>::subsistence_threshold();
 			let schedule = Contracts::current_schedule();
 			let mut ctx = MockContext::top_level(ALICE, &schedule);
+			let mut gas_meter = GasMeter::<Test>::new(GAS_LIMIT);
+			let executable = MockExecutable::from_storage(
+				rent_allowance_ch, &schedule, &mut gas_meter
+			).unwrap();
 			set_balance(&ALICE, subsistence * 10);
 
 			let result = ctx.instantiate(
 				subsistence * 5,
-				&mut GasMeter::<Test>::new(GAS_LIMIT),
-				MockExecutable::from_storage(rent_allowance_ch, &schedule).unwrap(),
+				&mut gas_meter,
+				executable,
 				vec![],
 				&[],
 			);

substrate/frame/contracts/src/gas.rs

@@ -15,14 +15,17 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::Config;
+use crate::{Config, Error};
 use sp_std::marker::PhantomData;
 use sp_runtime::traits::Zero;
 use frame_support::{
-	dispatch::{DispatchResultWithPostInfo, PostDispatchInfo, DispatchErrorWithPostInfo},
+	dispatch::{
+		DispatchResultWithPostInfo, PostDispatchInfo, DispatchErrorWithPostInfo, DispatchError,
+	},
 	weights::Weight,
 };
 use pallet_contracts_primitives::ExecError;
+use sp_core::crypto::UncheckedFrom;
 
 #[cfg(test)]
 use std::{any::Any, fmt::Debug};
@@ -30,22 +33,6 @@ use std::{any::Any, fmt::Debug};
 // Gas is essentially the same as weight. It is a 1 to 1 correspondence.
 pub type Gas = Weight;
 
-#[must_use]
-#[derive(Debug, PartialEq, Eq)]
-pub enum GasMeterResult {
-	Proceed(ChargedAmount),
-	OutOfGas,
-}
-
-impl GasMeterResult {
-	pub fn is_out_of_gas(&self) -> bool {
-		match *self {
-			GasMeterResult::OutOfGas => true,
-			GasMeterResult::Proceed(_) => false,
-		}
-	}
-}
-
 #[derive(Debug, PartialEq, Eq)]
 pub struct ChargedAmount(Gas);
 
@@ -103,7 +90,11 @@ pub struct GasMeter<T: Config> {
 	#[cfg(test)]
 	tokens: Vec<ErasedToken>,
 }
-impl<T: Config> GasMeter<T> {
+
+impl<T: Config> GasMeter<T>
+where
+	T::AccountId: UncheckedFrom<<T as frame_system::Config>::Hash> + AsRef<[u8]>
+{
 	pub fn new(gas_limit: Gas) -> Self {
 		GasMeter {
 			gas_limit,
@@ -128,7 +119,7 @@ impl<T: Config> GasMeter<T> {
 		&mut self,
 		metadata: &Tok::Metadata,
 		token: Tok,
-	) -> GasMeterResult {
+	) -> Result<ChargedAmount, DispatchError> {
 		#[cfg(test)]
 		{
 			// Unconditionally add the token to the storage.
@@ -149,11 +140,25 @@ impl<T: Config> GasMeter<T> {
 		self.gas_left = new_value.unwrap_or_else(Zero::zero);
 
 		match new_value {
-			Some(_) => GasMeterResult::Proceed(ChargedAmount(amount)),
-			None => GasMeterResult::OutOfGas,
+			Some(_) => Ok(ChargedAmount(amount)),
+			None => Err(Error::<T>::OutOfGas.into()),
 		}
 	}
 
+	/// Adjust a previously charged amount down to its actual amount.
+	///
+	/// This is when a maximum a priori amount was charged and then should be partially
+	/// refunded to match the actual amount.
+	pub fn adjust_gas<Tok: Token<T>>(
+		&mut self,
+		charged_amount: ChargedAmount,
+		metadata: &Tok::Metadata,
+		token: Tok,
+	) {
+		let adjustment = charged_amount.0.saturating_sub(token.calculate_amount(metadata));
+		self.gas_left = self.gas_left.saturating_add(adjustment).min(self.gas_limit);
+	}
+
 	/// Refund previously charged gas back to the gas meter.
 	///
 	/// This can be used if a gas worst case estimation must be charged before
@@ -304,7 +309,7 @@ mod tests {
 
 		let result = gas_meter
 			.charge(&MultiplierTokenMetadata { multiplier: 3 }, MultiplierToken(10));
-		assert!(!result.is_out_of_gas());
+		assert!(!result.is_err());
 
 		assert_eq!(gas_meter.gas_left(), 49_970);
 	}
@@ -312,10 +317,10 @@ mod tests {
 	#[test]
 	fn tracing() {
 		let mut gas_meter = GasMeter::<Test>::new(50000);
-		assert!(!gas_meter.charge(&(), SimpleToken(1)).is_out_of_gas());
+		assert!(!gas_meter.charge(&(), SimpleToken(1)).is_err());
 		assert!(!gas_meter
 			.charge(&MultiplierTokenMetadata { multiplier: 3 }, MultiplierToken(10))
-			.is_out_of_gas());
+			.is_err());
 
 		let mut tokens = gas_meter.tokens()[0..2].iter();
 		match_tokens!(tokens, SimpleToken(1), MultiplierToken(10),);
@@ -325,7 +330,7 @@ mod tests {
 	#[test]
 	fn refuse_to_execute_anything_if_zero() {
 		let mut gas_meter = GasMeter::<Test>::new(0);
-		assert!(gas_meter.charge(&(), SimpleToken(1)).is_out_of_gas());
+		assert!(gas_meter.charge(&(), SimpleToken(1)).is_err());
 	}
 
 	// Make sure that if the gas meter is charged by exceeding amount then not only an error
@@ -338,10 +343,10 @@ mod tests {
 		let mut gas_meter = GasMeter::<Test>::new(200);
 
 		// The first charge is should lead to OOG.
-		assert!(gas_meter.charge(&(), SimpleToken(300)).is_out_of_gas());
+		assert!(gas_meter.charge(&(), SimpleToken(300)).is_err());
 
 		// The gas meter is emptied at this moment, so this should also fail.
-		assert!(gas_meter.charge(&(), SimpleToken(1)).is_out_of_gas());
+		assert!(gas_meter.charge(&(), SimpleToken(1)).is_err());
 	}
 
 
@@ -350,6 +355,6 @@ mod tests {
 	#[test]
 	fn charge_exact_amount() {
 		let mut gas_meter = GasMeter::<Test>::new(25);
-		assert!(!gas_meter.charge(&(), SimpleToken(25)).is_out_of_gas());
+		assert!(!gas_meter.charge(&(), SimpleToken(25)).is_err());
 	}
 }

substrate/frame/contracts/src/lib.rs

@@ -80,7 +80,7 @@
 //! * [Balances](../pallet_balances/index.html)
 
 #![cfg_attr(not(feature = "std"), no_std)]
-#![cfg_attr(feature = "runtime-benchmarks", recursion_limit="256")]
+#![cfg_attr(feature = "runtime-benchmarks", recursion_limit="512")]
 
 #[macro_use]
 mod gas;
@@ -126,9 +126,9 @@ use frame_support::{
 };
 use frame_system::{ensure_signed, ensure_root, Module as System};
 use pallet_contracts_primitives::{
-	RentProjectionResult, GetStorageResult, ContractAccessError, ContractExecResult, ExecResult,
+	RentProjectionResult, GetStorageResult, ContractAccessError, ContractExecResult,
 };
-use frame_support::weights::Weight;
+use frame_support::weights::{Weight, PostDispatchInfo, WithPostDispatchInfo};
 
 pub type CodeHash<T> = <T as frame_system::Config>::Hash;
 pub type TrieId = Vec<u8>;
@@ -344,6 +344,11 @@ pub trait Config: frame_system::Config {
 
 	/// The maximum amount of weight that can be consumed per block for lazy trie removal.
 	type DeletionWeightLimit: Get<Weight>;
+
+	/// The maximum length of a contract code in bytes. This limit applies to the instrumented
+	/// version of the code. Therefore `instantiate_with_code` can fail even when supplying
+	/// a wasm binary below this maximum size.
+	type MaxCodeSize: Get<u32>;
 }
 
 decl_error! {
@@ -538,7 +543,7 @@ decl_module! {
 		/// * If the account is a regular account, any value will be transferred.
 		/// * If no account exists and the call value is not less than `existential_deposit`,
 		/// a regular account will be created and any value will be transferred.
-		#[weight = T::WeightInfo::call().saturating_add(*gas_limit)]
+		#[weight = T::WeightInfo::call(T::MaxCodeSize::get() / 1024).saturating_add(*gas_limit)]
 		pub fn call(
 			origin,
 			dest: <T::Lookup as StaticLookup>::Source,
@@ -549,10 +554,13 @@ 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, T::WeightInfo::call())
+			let schedule = <Module<T>>::current_schedule();
+			let mut ctx = ExecutionContext::<T, PrefabWasmModule<T>>::top_level(origin, &schedule);
+			let (result, code_len) = match ctx.call(dest, value, &mut gas_meter, data) {
+				Ok((output, len)) => (Ok(output), len),
+				Err((err, len)) => (Err(err), len),
+			};
+			gas_meter.into_dispatch_result(result, T::WeightInfo::call(code_len / 1024))
 		}
 
 		/// Instantiates a new contract from the supplied `code` optionally transferring
@@ -592,16 +600,16 @@ decl_module! {
 			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);
+			ensure!(code_len <= T::MaxCodeSize::get(), 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)
-			});
+			let schedule = <Module<T>>::current_schedule();
+			let executable = PrefabWasmModule::from_code(code, &schedule)?;
+			let code_len = executable.code_len();
+			ensure!(code_len <= T::MaxCodeSize::get(), Error::<T>::CodeTooLarge);
+			let mut ctx = ExecutionContext::<T, PrefabWasmModule<T>>::top_level(origin, &schedule);
+			let result = ctx.instantiate(endowment, &mut gas_meter, executable, data, &salt)
+				.map(|(_address, output)| output);
 			gas_meter.into_dispatch_result(
 				result,
 				T::WeightInfo::instantiate_with_code(code_len / 1024, salt.len() as u32 / 1024)
@@ -614,8 +622,8 @@ decl_module! {
 		/// 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)
+			T::WeightInfo::instantiate(T::MaxCodeSize::get() / 1024, salt.len() as u32 / 1024)
+				.saturating_add(*gas_limit)
 		]
 		pub fn instantiate(
 			origin,
@@ -627,15 +635,15 @@ 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| {
-				let executable = PrefabWasmModule::from_storage(code_hash, &ctx.schedule)?;
-				let result = ctx.instantiate(endowment, gas_meter, executable, data, &salt)
-					.map(|(_address, output)| output)?;
-				Ok(result)
-			});
+			let schedule = <Module<T>>::current_schedule();
+			let executable = PrefabWasmModule::from_storage(code_hash, &schedule, &mut gas_meter)?;
+			let mut ctx = ExecutionContext::<T, PrefabWasmModule<T>>::top_level(origin, &schedule);
+			let code_len = executable.code_len();
+			let result = ctx.instantiate(endowment, &mut gas_meter, executable, data, &salt)
+				.map(|(_address, output)| output);
 			gas_meter.into_dispatch_result(
 				result,
-				T::WeightInfo::instantiate(salt.len() as u32 / 1024)
+				T::WeightInfo::instantiate(code_len / 1024, salt.len() as u32 / 1024),
 			)
 		}
 
@@ -648,7 +656,7 @@ decl_module! {
 		///
 		/// If contract is not evicted as a result of this call, [`Error::ContractNotEvictable`]
 		/// is returned and the sender is not eligible for the reward.
-		#[weight = T::WeightInfo::claim_surcharge()]
+		#[weight = T::WeightInfo::claim_surcharge(T::MaxCodeSize::get() / 1024)]
 		pub fn claim_surcharge(
 			origin,
 			dest: T::AccountId,
@@ -675,23 +683,26 @@ decl_module! {
 			};
 
 			// If poking the contract has lead to eviction of the contract, give out the rewards.
-			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),
-				)
-				.map(|_| Pays::No.into())
-				.map_err(Into::into)
-			} else {
-				Err(Error::<T>::ContractNotEvictable.into())
+			match Rent::<T, PrefabWasmModule<T>>::try_eviction(&dest, handicap)? {
+				(Some(rent_payed), code_len) => {
+					T::Currency::deposit_into_existing(
+						&rewarded,
+						T::SurchargeReward::get().min(rent_payed),
+					)
+					.map(|_| PostDispatchInfo {
+						actual_weight: Some(T::WeightInfo::claim_surcharge(code_len / 1024)),
+						pays_fee: Pays::No,
+					})
+					.map_err(Into::into)
+				}
+				(None, code_len) => Err(Error::<T>::ContractNotEvictable.with_weight(
+					T::WeightInfo::claim_surcharge(code_len / 1024)
+				)),
 			}
 		}
 	}
 }
 
-/// Public APIs provided by the contracts module.
 impl<T: Config> Module<T>
 where
 	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>,
@@ -710,12 +721,12 @@ where
 		input_data: Vec<u8>,
 	) -> ContractExecResult {
 		let mut gas_meter = GasMeter::new(gas_limit);
-		let exec_result = Self::execute_wasm(origin, &mut gas_meter, |ctx, gas_meter| {
-			ctx.call(dest, value, gas_meter, input_data)
-		});
+		let schedule = <Module<T>>::current_schedule();
+		let mut ctx = ExecutionContext::<T, PrefabWasmModule<T>>::top_level(origin, &schedule);
+		let result = ctx.call(dest, value, &mut gas_meter, input_data);
 		let gas_consumed = gas_meter.gas_spent();
 		ContractExecResult {
-			exec_result,
+			exec_result: result.map(|r| r.0).map_err(|r| r.0),
 			gas_consumed,
 		}
 	}
@@ -731,18 +742,12 @@ where
 		Ok(maybe_value)
 	}
 
+	/// Query how many blocks the contract stays alive given that the amount endowment
+	/// and consumed storage does not change.
 	pub fn rent_projection(address: T::AccountId) -> RentProjectionResult<T::BlockNumber> {
 		Rent::<T, PrefabWasmModule<T>>::compute_projection(&address)
 	}
 
-	/// Store code for benchmarks which does not check nor instrument the code.
-	#[cfg(feature = "runtime-benchmarks")]
-	pub fn store_code_raw(code: Vec<u8>) -> DispatchResult {
-		let schedule = <Module<T>>::current_schedule();
-		PrefabWasmModule::store_code_unchecked(code, &schedule)?;
-		Ok(())
-	}
-
 	/// Determine the address of a contract,
 	///
 	/// This is the address generation function used by contract instantiation. Its result
@@ -775,23 +780,22 @@ where
 	pub fn subsistence_threshold() -> BalanceOf<T> {
 		T::Currency::minimum_balance().saturating_add(T::TombstoneDeposit::get())
 	}
-}
 
-impl<T: Config> Module<T>
-where
-	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>,
-{
-	fn execute_wasm(
-		origin: T::AccountId,
-		gas_meter: &mut GasMeter<T>,
-		func: impl FnOnce(
-			&mut ExecutionContext<T, PrefabWasmModule<T>>,
-			&mut GasMeter<T>,
-		) -> ExecResult,
-	) -> ExecResult {
+	/// Store code for benchmarks which does not check nor instrument the code.
+	#[cfg(feature = "runtime-benchmarks")]
+	fn store_code_raw(code: Vec<u8>) -> DispatchResult {
 		let schedule = <Module<T>>::current_schedule();
-		let mut ctx = ExecutionContext::top_level(origin, &schedule);
-		func(&mut ctx, gas_meter)
+		PrefabWasmModule::store_code_unchecked(code, &schedule)?;
+		Ok(())
+	}
+
+	/// This exists so that benchmarks can determine the weight of running an instrumentation.
+	#[cfg(feature = "runtime-benchmarks")]
+	fn reinstrument_module(
+		module: &mut PrefabWasmModule<T>,
+		schedule: &Schedule<T>
+	) -> DispatchResult {
+		self::wasm::reinstrument(module, schedule)
 	}
 }
 

substrate/frame/contracts/src/rent.rs

@@ -325,13 +325,14 @@ where
 	pub fn try_eviction(
 		account: &T::AccountId,
 		handicap: T::BlockNumber,
-	) -> Result<Option<BalanceOf<T>>, DispatchError> {
+	) -> Result<(Option<BalanceOf<T>>, u32), DispatchError> {
 		let contract = <ContractInfoOf<T>>::get(account);
 		let contract = match contract {
-			None | Some(ContractInfo::Tombstone(_)) => return Ok(None),
+			None | Some(ContractInfo::Tombstone(_)) => return Ok((None, 0)),
 			Some(ContractInfo::Alive(contract)) => contract,
 		};
 		let module = PrefabWasmModule::<T>::from_storage_noinstr(contract.code_hash)?;
+		let code_len = module.code_len();
 		let current_block_number = <frame_system::Module<T>>::block_number();
 		let verdict = Self::consider_case(
 			account,
@@ -353,9 +354,9 @@ where
 				Self::enact_verdict(
 					account, contract, current_block_number, verdict, Some(module),
 				)?;
-				Ok(Some(rent_payed))
+				Ok((Some(rent_payed), code_len))
 			}
-			_ => Ok(None),
+			_ => Ok((None, code_len)),
 		}
 	}
 
@@ -447,28 +448,32 @@ where
 	/// Upon succesful restoration, `origin` will be destroyed, all its funds are transferred to
 	/// the restored account. The restored account will inherit the last write block and its last
 	/// deduct block will be set to the current block.
+	///
+	/// # Return Value
+	///
+	/// Result<(CallerCodeSize, DestCodeSize), (DispatchError, CallerCodeSize, DestCodesize)>
 	pub fn restore_to(
 		origin: T::AccountId,
 		dest: T::AccountId,
 		code_hash: CodeHash<T>,
 		rent_allowance: BalanceOf<T>,
 		delta: Vec<crate::exec::StorageKey>,
-	) -> Result<(), DispatchError> {
+	) -> Result<(u32, u32), (DispatchError, u32, u32)> {
 		let mut origin_contract = <ContractInfoOf<T>>::get(&origin)
 			.and_then(|c| c.get_alive())
-			.ok_or(Error::<T>::InvalidSourceContract)?;
+			.ok_or((Error::<T>::InvalidSourceContract.into(), 0, 0))?;
 
 		let child_trie_info = origin_contract.child_trie_info();
 
 		let current_block = <frame_system::Module<T>>::block_number();
 
 		if origin_contract.last_write == Some(current_block) {
-			return Err(Error::<T>::InvalidContractOrigin.into());
+			return Err((Error::<T>::InvalidContractOrigin.into(), 0, 0));
 		}
 
 		let dest_tombstone = <ContractInfoOf<T>>::get(&dest)
 			.and_then(|c| c.get_tombstone())
-			.ok_or(Error::<T>::InvalidDestinationContract)?;
+			.ok_or((Error::<T>::InvalidDestinationContract.into(), 0, 0))?;
 
 		let last_write = if !delta.is_empty() {
 			Some(current_block)
@@ -477,7 +482,7 @@ where
 		};
 
 		// Fails if the code hash does not exist on chain
-		E::add_user(code_hash)?;
+		let caller_code_len = E::add_user(code_hash).map_err(|e| (e, 0, 0))?;
 
 		// We are allowed to eagerly modify storage even though the function can
 		// fail later due to tombstones not matching. This is because the restoration
@@ -501,13 +506,13 @@ where
 		);
 
 		if tombstone != dest_tombstone {
-			return Err(Error::<T>::InvalidTombstone.into());
+			return Err((Error::<T>::InvalidTombstone.into(), caller_code_len, 0));
 		}
 
 		origin_contract.storage_size -= bytes_taken;
 
 		<ContractInfoOf<T>>::remove(&origin);
-		E::remove_user(origin_contract.code_hash);
+		let tombstone_code_len = 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,
@@ -523,6 +528,6 @@ where
 		T::Currency::make_free_balance_be(&origin, <BalanceOf<T>>::zero());
 		T::Currency::deposit_creating(&dest, origin_free_balance);
 
-		Ok(())
+		Ok((caller_code_len, tombstone_code_len))
 	}
 }

substrate/frame/contracts/src/schedule.rs

@@ -104,10 +104,6 @@ pub struct Limits {
 
 	/// The maximum length of a subject in bytes used for PRNG generation.
 	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 `instantiate_with_code`.
-	pub code_size: u32,
 }
 
 impl Limits {
@@ -250,9 +246,18 @@ pub struct HostFnWeights<T: Config> {
 	/// Weight of calling `seal_terminate`.
 	pub terminate: Weight,
 
+	/// Weight per byte of the terminated contract.
+	pub terminate_per_code_byte: Weight,
+
 	/// Weight of calling `seal_restore_to`.
 	pub restore_to: Weight,
 
+	/// Weight per byte of the restoring contract.
+	pub restore_to_per_caller_code_byte: Weight,
+
+	/// Weight per byte of the restored contract.
+	pub restore_to_per_tombstone_code_byte: Weight,
+
 	/// Weight per delta key supplied to `seal_restore_to`.
 	pub restore_to_per_delta: Weight,
 
@@ -292,6 +297,9 @@ pub struct HostFnWeights<T: Config> {
 	/// Weight of calling `seal_call`.
 	pub call: Weight,
 
+	/// Weight per byte of the called contract.
+	pub call_per_code_byte: Weight,
+
 	/// Weight surcharge that is claimed if `seal_call` does a balance transfer.
 	pub call_transfer_surcharge: Weight,
 
@@ -304,6 +312,9 @@ pub struct HostFnWeights<T: Config> {
 	/// Weight of calling `seal_instantiate`.
 	pub instantiate: Weight,
 
+	/// Weight per byte of the instantiated contract.
+	pub instantiate_per_code_byte: Weight,
+
 	/// Weight per input byte supplied to `seal_instantiate`.
 	pub instantiate_per_input_byte: Weight,
 
@@ -443,7 +454,6 @@ impl Default for Limits {
 			table_size: 4096,
 			br_table_size: 256,
 			subject_len: 32,
-			code_size: 512 * 1024,
 		}
 	}
 }
@@ -528,8 +538,11 @@ impl<T: Config> Default for HostFnWeights<T> {
 			r#return: cost!(seal_return),
 			return_per_byte: cost_byte!(seal_return_per_kb),
 			terminate: cost!(seal_terminate),
+			terminate_per_code_byte: cost_byte!(seal_terminate_per_code_kb),
 			restore_to: cost!(seal_restore_to),
-			restore_to_per_delta: cost_batched!(seal_restore_to_per_delta),
+			restore_to_per_caller_code_byte: cost_byte_args!(seal_restore_to_per_code_kb_delta, 1, 0, 0),
+			restore_to_per_tombstone_code_byte: cost_byte_args!(seal_restore_to_per_code_kb_delta, 0, 1, 0),
+			restore_to_per_delta: cost_batched_args!(seal_restore_to_per_code_kb_delta, 0, 0, 1),
 			random: cost_batched!(seal_random),
 			deposit_event: cost_batched!(seal_deposit_event),
 			deposit_event_per_topic: cost_batched_args!(seal_deposit_event_per_topic_and_kb, 1, 0),
@@ -542,13 +555,15 @@ impl<T: Config> Default for HostFnWeights<T> {
 			get_storage_per_byte: cost_byte_batched!(seal_get_storage_per_kb),
 			transfer: cost_batched!(seal_transfer),
 			call: cost_batched!(seal_call),
-			call_transfer_surcharge: cost_batched_args!(seal_call_per_transfer_input_output_kb, 1, 0, 0),
-			call_per_input_byte: cost_byte_batched_args!(seal_call_per_transfer_input_output_kb, 0, 1, 0),
-			call_per_output_byte: cost_byte_batched_args!(seal_call_per_transfer_input_output_kb, 0, 0, 1),
+			call_per_code_byte: cost_byte_batched_args!(seal_call_per_code_transfer_input_output_kb, 1, 0, 0, 0),
+			call_transfer_surcharge: cost_batched_args!(seal_call_per_code_transfer_input_output_kb, 0, 1, 0, 0),
+			call_per_input_byte: cost_byte_batched_args!(seal_call_per_code_transfer_input_output_kb, 0, 0, 1, 0),
+			call_per_output_byte: cost_byte_batched_args!(seal_call_per_code_transfer_input_output_kb, 0, 0, 0, 1),
 			instantiate: cost_batched!(seal_instantiate),
-			instantiate_per_input_byte: cost_byte_batched_args!(seal_instantiate_per_input_output_salt_kb, 1, 0, 0),
-			instantiate_per_output_byte: cost_byte_batched_args!(seal_instantiate_per_input_output_salt_kb, 0, 1, 0),
-			instantiate_per_salt_byte: cost_byte_batched_args!(seal_instantiate_per_input_output_salt_kb, 0, 0, 1),
+			instantiate_per_code_byte: cost_byte_batched_args!(seal_instantiate_per_code_input_output_salt_kb, 1, 0, 0, 0),
+			instantiate_per_input_byte: cost_byte_batched_args!(seal_instantiate_per_code_input_output_salt_kb, 0, 1, 0, 0),
+			instantiate_per_output_byte: cost_byte_batched_args!(seal_instantiate_per_code_input_output_salt_kb, 0, 0, 1, 0),
+			instantiate_per_salt_byte: cost_byte_batched_args!(seal_instantiate_per_code_input_output_salt_kb, 0, 0, 0, 1),
 			hash_sha2_256: cost_batched!(seal_hash_sha2_256),
 			hash_sha2_256_per_byte: cost_byte_batched!(seal_hash_sha2_256_per_kb),
 			hash_keccak_256: cost_batched!(seal_hash_keccak_256),

substrate/frame/contracts/src/tests.rs

@@ -24,6 +24,7 @@ use crate::{
 		UncheckedFrom, InitState, ReturnFlags,
 	},
 	exec::{AccountIdOf, Executable}, wasm::PrefabWasmModule,
+	weights::WeightInfo,
 };
 use assert_matches::assert_matches;
 use codec::Encode;
@@ -35,7 +36,7 @@ use sp_runtime::{
 use sp_io::hashing::blake2_256;
 use frame_support::{
 	assert_ok, assert_err, assert_err_ignore_postinfo,
-	parameter_types, StorageMap, assert_storage_noop,
+	parameter_types, StorageMap, StorageValue, assert_storage_noop,
 	traits::{Currency, ReservableCurrency, OnInitialize},
 	weights::{Weight, PostDispatchInfo, DispatchClass, constants::WEIGHT_PER_SECOND},
 	dispatch::DispatchErrorWithPostInfo,
@@ -250,6 +251,7 @@ parameter_types! {
 	pub const MaxValueSize: u32 = 16_384;
 	pub const DeletionQueueDepth: u32 = 1024;
 	pub const DeletionWeightLimit: Weight = 500_000_000_000;
+	pub const MaxCodeSize: u32 = 2 * 1024;
 }
 
 parameter_types! {
@@ -282,6 +284,7 @@ impl Config for Test {
 	type ChainExtension = TestExtension;
 	type DeletionQueueDepth = DeletionQueueDepth;
 	type DeletionWeightLimit = DeletionWeightLimit;
+	type MaxCodeSize = MaxCodeSize;
 }
 
 pub const ALICE: AccountId32 = AccountId32::new([1u8; 32]);
@@ -350,7 +353,7 @@ where
 fn calling_plain_account_fails() {
 	ExtBuilder::default().build().execute_with(|| {
 		let _ = Balances::deposit_creating(&ALICE, 100_000_000);
-		let base_cost = <<Test as crate::Config>::WeightInfo as crate::WeightInfo>::call();
+		let base_cost = <<Test as Config>::WeightInfo as WeightInfo>::call(0);
 
 		assert_eq!(
 			Contracts::call(Origin::signed(ALICE), BOB, 0, GAS_LIMIT, Vec::new()),
@@ -2432,7 +2435,7 @@ fn lazy_removal_does_no_run_on_full_block() {
 		// Run the lazy removal without any limit so that all keys would be removed if there
 		// had been some weight left in the block.
 		let weight_used = Contracts::on_initialize(Weight::max_value());
-		let base = <<Test as crate::Config>::WeightInfo as crate::WeightInfo>::on_initialize();
+		let base = <<Test as Config>::WeightInfo as WeightInfo>::on_initialize();
 		assert_eq!(weight_used, base);
 
 		// All the keys are still in place
@@ -2717,3 +2720,69 @@ fn refcounter() {
 		assert_matches!(crate::CodeStorage::<Test>::get(code_hash), None);
 	});
 }
+
+
+#[test]
+fn reinstrument_does_charge() {
+	let (wasm, code_hash) = compile_module::<Test>("return_with_data").unwrap();
+	ExtBuilder::default().existential_deposit(50).build().execute_with(|| {
+		let _ = Balances::deposit_creating(&ALICE, 1_000_000);
+		let subsistence = Module::<Test>::subsistence_threshold();
+		let zero = 0u32.to_le_bytes().encode();
+		let code_len = wasm.len() as u32;
+
+		assert_ok!(Contracts::instantiate_with_code(
+			Origin::signed(ALICE),
+			subsistence * 100,
+			GAS_LIMIT,
+			wasm,
+			zero.clone(),
+			vec![],
+		));
+
+		let addr = Contracts::contract_address(&ALICE, &code_hash, &[]);
+
+		// Call the contract two times without reinstrument
+
+		let result0 = Contracts::bare_call(
+			ALICE,
+			addr.clone(),
+			0,
+			GAS_LIMIT,
+			zero.clone(),
+		);
+		assert!(result0.exec_result.unwrap().is_success());
+
+		let result1 = Contracts::bare_call(
+			ALICE,
+			addr.clone(),
+			0,
+			GAS_LIMIT,
+			zero.clone(),
+		);
+		assert!(result1.exec_result.unwrap().is_success());
+
+		// They should match because both where called with the same schedule.
+		assert_eq!(result0.gas_consumed, result1.gas_consumed);
+
+		// Update the schedule version but keep the rest the same
+		crate::CurrentSchedule::mutate(|old: &mut Schedule<Test>| {
+			old.version += 1;
+		});
+
+		// This call should trigger reinstrumentation
+		let result2 = Contracts::bare_call(
+			ALICE,
+			addr.clone(),
+			0,
+			GAS_LIMIT,
+			zero.clone(),
+		);
+		assert!(result2.exec_result.unwrap().is_success());
+		assert!(result2.gas_consumed > result1.gas_consumed);
+		assert_eq!(
+			result2.gas_consumed,
+			result1.gas_consumed + <Test as Config>::WeightInfo::instrument(code_len / 1024),
+		);
+	});
+}

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

@@ -30,9 +30,13 @@
 use crate::{
 	CodeHash, CodeStorage, PristineCode, Schedule, Config, Error,
 	wasm::{prepare, PrefabWasmModule}, Module as Contracts, RawEvent,
+	gas::{Gas, GasMeter, Token},
+	weights::WeightInfo,
 };
 use sp_core::crypto::UncheckedFrom;
-use frame_support::{StorageMap, dispatch::{DispatchError, DispatchResult}};
+use frame_support::{StorageMap, dispatch::DispatchError};
+#[cfg(feature = "runtime-benchmarks")]
+pub use self::private::reinstrument as reinstrument;
 
 /// Put the instrumented module in storage.
 ///
@@ -77,14 +81,14 @@ where
 }
 
 /// Increment the refcount of a code in-storage by one.
-pub fn increment_refcount<T: Config>(code_hash: CodeHash<T>) -> DispatchResult
+pub fn increment_refcount<T: Config>(code_hash: CodeHash<T>) -> Result<u32, DispatchError>
 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(())
+			Ok(module.original_code_len)
 		} else {
 			Err(Error::<T>::CodeNotFound.into())
 		}
@@ -92,19 +96,23 @@ where
 }
 
 /// 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>)
+pub fn decrement_refcount<T: Config>(code_hash: CodeHash<T>) -> u32
 where
 	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>
 {
 	<CodeStorage<T>>::mutate_exists(code_hash, |existing| {
 		if let Some(module) = existing {
+			let code_len = module.original_code_len;
 			module.refcount = module.refcount.saturating_sub(1);
 			if module.refcount == 0 {
 				*existing = None;
 				finish_removal::<T>(code_hash);
 			}
+			code_len
+		} else {
+			0
 		}
-	});
+	})
 }
 
 /// Load code with the given code hash.
@@ -114,31 +122,48 @@ where
 /// re-instrumentation and update the cache in the storage.
 pub fn load<T: Config>(
 	code_hash: CodeHash<T>,
-	schedule: Option<&Schedule<T>>,
+	reinstrument: Option<(&Schedule<T>, &mut GasMeter<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)?;
+	prefab_module.code_hash = code_hash;
 
-	if let Some(schedule) = schedule {
+	if let Some((schedule, gas_meter)) = reinstrument {
 		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);
+			gas_meter.charge(&(), InstrumentToken(prefab_module.original_code_len))?;
+			private::reinstrument(&mut prefab_module, schedule)?;
 		}
 	}
-	prefab_module.code_hash = code_hash;
 	Ok(prefab_module)
 }
 
+mod private {
+	use super::*;
+
+	/// Instruments the passed prefab wasm module with the supplied schedule.
+	pub fn reinstrument<T: Config>(
+		prefab_module: &mut PrefabWasmModule<T>,
+		schedule: &Schedule<T>,
+	) -> Result<(), DispatchError>
+	where
+		T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>
+	{
+		let original_code = <PristineCode<T>>::get(&prefab_module.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(&prefab_module.code_hash, &*prefab_module);
+		Ok(())
+	}
+}
+
 /// Finish removal of a code by deleting the pristine code and emitting an event.
 fn finish_removal<T: Config>(code_hash: CodeHash<T>)
 where
@@ -161,3 +186,17 @@ fn increment_64(refcount: &mut u64) {
 		qed
 	");
 }
+
+/// Token to be supplied to the gas meter which charges the weight needed for reinstrumenting
+/// a contract of the specified size in bytes.
+#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
+#[derive(Clone, Copy)]
+struct InstrumentToken(u32);
+
+impl<T: Config> Token<T> for InstrumentToken {
+	type Metadata = ();
+
+	fn calculate_amount(&self, _metadata: &Self::Metadata) -> Gas {
+		T::WeightInfo::instrument(self.0 / 1024)
+	}
+}

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

@@ -33,9 +33,11 @@ use crate::{
 use sp_std::prelude::*;
 use sp_core::crypto::UncheckedFrom;
 use codec::{Encode, Decode};
-use frame_support::dispatch::{DispatchError, DispatchResult};
+use frame_support::dispatch::DispatchError;
 use pallet_contracts_primitives::ExecResult;
 pub use self::runtime::{ReturnCode, Runtime, RuntimeToken};
+#[cfg(feature = "runtime-benchmarks")]
+pub use self::code_cache::reinstrument;
 
 /// A prepared wasm module ready for execution.
 ///
@@ -125,7 +127,7 @@ where
 	pub fn store_code_unchecked(
 		original_code: Vec<u8>,
 		schedule: &Schedule<T>
-	) -> DispatchResult {
+	) -> Result<(), DispatchError> {
 		let executable = prepare::benchmarking::prepare_contract(original_code, schedule)
 			.map_err::<DispatchError, _>(Into::into)?;
 		code_cache::store(executable);
@@ -145,9 +147,10 @@ where
 {
 	fn from_storage(
 		code_hash: CodeHash<T>,
-		schedule: &Schedule<T>
+		schedule: &Schedule<T>,
+		gas_meter: &mut GasMeter<T>,
 	) -> Result<Self, DispatchError> {
-		code_cache::load(code_hash, Some(schedule))
+		code_cache::load(code_hash, Some((schedule, gas_meter)))
 	}
 
 	fn from_storage_noinstr(code_hash: CodeHash<T>) -> Result<Self, DispatchError> {
@@ -158,11 +161,11 @@ where
 		code_cache::store_decremented(self);
 	}
 
-	fn add_user(code_hash: CodeHash<T>) -> DispatchResult {
+	fn add_user(code_hash: CodeHash<T>) -> Result<u32, DispatchError> {
 		code_cache::increment_refcount::<T>(code_hash)
 	}
 
-	fn remove_user(code_hash: CodeHash<T>) {
+	fn remove_user(code_hash: CodeHash<T>) -> u32 {
 		code_cache::decrement_refcount::<T>(code_hash)
 	}
 
@@ -222,6 +225,10 @@ where
 		let len = self.original_code_len.saturating_add(self.code.len() as u32);
 		len.checked_div(self.refcount as u32).unwrap_or(len)
 	}
+
+	fn code_len(&self) -> u32 {
+		self.code.len() as u32
+	}
 }
 
 #[cfg(test)]
@@ -305,7 +312,7 @@ mod tests {
 			gas_meter: &mut GasMeter<Test>,
 			data: Vec<u8>,
 			salt: &[u8],
-		) -> Result<(AccountIdOf<Self::T>, ExecReturnValue), ExecError> {
+		) -> Result<(AccountIdOf<Self::T>, ExecReturnValue, u32), (ExecError, u32)> {
 			self.instantiates.push(InstantiateEntry {
 				code_hash: code_hash.clone(),
 				endowment,
@@ -319,6 +326,7 @@ mod tests {
 					flags: ReturnFlags::empty(),
 					data: Vec::new(),
 				},
+				0,
 			))
 		}
 		fn transfer(
@@ -339,7 +347,7 @@ mod tests {
 			value: u64,
 			_gas_meter: &mut GasMeter<Test>,
 			data: Vec<u8>,
-		) -> ExecResult {
+		) -> Result<(ExecReturnValue, u32), (ExecError, u32)> {
 			self.transfers.push(TransferEntry {
 				to: to.clone(),
 				value,
@@ -347,16 +355,16 @@ mod tests {
 			});
 			// Assume for now that it was just a plain transfer.
 			// TODO: Add tests for different call outcomes.
-			Ok(ExecReturnValue { flags: ReturnFlags::empty(), data: Vec::new() })
+			Ok((ExecReturnValue { flags: ReturnFlags::empty(), data: Vec::new() }, 0))
 		}
 		fn terminate(
 			&mut self,
 			beneficiary: &AccountIdOf<Self::T>,
-		) -> Result<(), DispatchError> {
+		) -> Result<u32, (DispatchError, u32)> {
 			self.terminations.push(TerminationEntry {
 				beneficiary: beneficiary.clone(),
 			});
-			Ok(())
+			Ok(0)
 		}
 		fn restore_to(
 			&mut self,
@@ -364,14 +372,14 @@ mod tests {
 			code_hash: H256,
 			rent_allowance: u64,
 			delta: Vec<StorageKey>,
-		) -> Result<(), DispatchError> {
+		) -> Result<(u32, u32), (DispatchError, u32, u32)> {
 			self.restores.push(RestoreEntry {
 				dest,
 				code_hash,
 				rent_allowance,
 				delta,
 			});
-			Ok(())
+			Ok((0, 0))
 		}
 		fn caller(&self) -> &AccountIdOf<Self::T> {
 			&ALICE
@@ -443,7 +451,7 @@ mod tests {
 			gas_meter: &mut GasMeter<Test>,
 			input_data: Vec<u8>,
 			salt: &[u8],
-		) -> Result<(AccountIdOf<Self::T>, ExecReturnValue), ExecError> {
+		) -> Result<(AccountIdOf<Self::T>, ExecReturnValue, u32), (ExecError, u32)> {
 			(**self).instantiate(code, value, gas_meter, input_data, salt)
 		}
 		fn transfer(
@@ -456,7 +464,7 @@ mod tests {
 		fn terminate(
 			&mut self,
 			beneficiary: &AccountIdOf<Self::T>,
-		) -> Result<(), DispatchError> {
+		) -> Result<u32, (DispatchError, u32)> {
 			(**self).terminate(beneficiary)
 		}
 		fn call(
@@ -465,7 +473,7 @@ mod tests {
 			value: u64,
 			gas_meter: &mut GasMeter<Test>,
 			input_data: Vec<u8>,
-		) -> ExecResult {
+		) -> Result<(ExecReturnValue, u32), (ExecError, u32)> {
 			(**self).call(to, value, gas_meter, input_data)
 		}
 		fn restore_to(
@@ -474,7 +482,7 @@ mod tests {
 			code_hash: H256,
 			rent_allowance: u64,
 			delta: Vec<StorageKey>,
-		) -> Result<(), DispatchError> {
+		) -> Result<(u32, u32), (DispatchError, u32, u32)> {
 			(**self).restore_to(
 				dest,
 				code_hash,

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

@@ -20,11 +20,11 @@
 use crate::{
 	HostFnWeights, Config, CodeHash, BalanceOf, Error,
 	exec::{Ext, StorageKey, TopicOf},
-	gas::{Gas, GasMeter, Token, GasMeterResult, ChargedAmount},
+	gas::{Gas, GasMeter, Token, ChargedAmount},
 	wasm::env_def::ConvertibleToWasm,
 };
 use parity_wasm::elements::ValueType;
-use frame_support::{dispatch::DispatchError, ensure};
+use frame_support::{dispatch::DispatchError, ensure, traits::Get};
 use sp_std::prelude::*;
 use codec::{Decode, DecodeAll, Encode};
 use sp_runtime::traits::SaturatedConversion;
@@ -165,8 +165,12 @@ pub enum RuntimeToken {
 	Return(u32),
 	/// Weight of calling `seal_terminate`.
 	Terminate,
+	/// Weight that is added to `seal_terminate` for every byte of the terminated contract.
+	TerminateSurchargeCodeSize(u32),
 	/// Weight of calling `seal_restore_to` per number of supplied delta entries.
 	RestoreTo(u32),
+	/// Weight that is added to `seal_restore_to` for the involved code sizes.
+	RestoreToSurchargeCodeSize{caller_code: u32, tombstone_code: u32},
 	/// Weight of calling `seal_random`. It includes the weight for copying the subject.
 	Random,
 	/// Weight of calling `seal_reposit_event` with the given number of topics and event size.
@@ -185,6 +189,8 @@ pub enum RuntimeToken {
 	Transfer,
 	/// Weight of calling `seal_call` for the given input size.
 	CallBase(u32),
+	/// Weight that is added to `seal_call` for every byte of the called contract.
+	CallSurchargeCodeSize(u32),
 	/// Weight of the transfer performed during a call.
 	CallSurchargeTransfer,
 	/// Weight of output received through `seal_call` for the given size.
@@ -193,6 +199,8 @@ pub enum RuntimeToken {
 	/// This includes the transfer as an instantiate without a value will always be below
 	/// the existential deposit and is disregarded as corner case.
 	InstantiateBase{input_data_len: u32, salt_len: u32},
+	/// Weight that is added to `seal_instantiate` for every byte of the instantiated contract.
+	InstantiateSurchargeCodeSize(u32),
 	/// Weight of output received through `seal_instantiate` for the given size.
 	InstantiateCopyOut(u32),
 	/// Weight of calling `seal_hash_sha_256` for the given input size.
@@ -235,8 +243,13 @@ where
 			Return(len) => s.r#return
 				.saturating_add(s.return_per_byte.saturating_mul(len.into())),
 			Terminate => s.terminate,
+			TerminateSurchargeCodeSize(len) => s.terminate_per_code_byte.saturating_mul(len.into()),
 			RestoreTo(delta) => s.restore_to
 				.saturating_add(s.restore_to_per_delta.saturating_mul(delta.into())),
+			RestoreToSurchargeCodeSize{caller_code, tombstone_code} =>
+				s.restore_to_per_caller_code_byte.saturating_mul(caller_code.into()).saturating_add(
+					s.restore_to_per_tombstone_code_byte.saturating_mul(tombstone_code.into())
+				),
 			Random => s.random,
 			DepositEvent{num_topic, len} => s.deposit_event
 				.saturating_add(s.deposit_event_per_topic.saturating_mul(num_topic.into()))
@@ -250,11 +263,14 @@ where
 			Transfer => s.transfer,
 			CallBase(len) => s.call
 				.saturating_add(s.call_per_input_byte.saturating_mul(len.into())),
+			CallSurchargeCodeSize(len) => s.call_per_code_byte.saturating_mul(len.into()),
 			CallSurchargeTransfer => s.call_transfer_surcharge,
 			CallCopyOut(len) => s.call_per_output_byte.saturating_mul(len.into()),
 			InstantiateBase{input_data_len, salt_len} => s.instantiate
 				.saturating_add(s.instantiate_per_input_byte.saturating_mul(input_data_len.into()))
 				.saturating_add(s.instantiate_per_salt_byte.saturating_mul(salt_len.into())),
+			InstantiateSurchargeCodeSize(len) =>
+				s.instantiate_per_code_byte.saturating_mul(len.into()),
 			InstantiateCopyOut(len) => s.instantiate_per_output_byte
 				.saturating_mul(len.into()),
 			HashSha256(len) => s.hash_sha2_256
@@ -408,10 +424,19 @@ where
 	where
 		Tok: Token<E::T, Metadata=HostFnWeights<E::T>>,
 	{
-		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())
-		}
+		self.gas_meter.charge(&self.ext.schedule().host_fn_weights, token)
+	}
+
+	/// Correct previously charged gas amount.
+	pub fn adjust_gas<Tok>(&mut self, charged_amount: ChargedAmount, adjusted_amount: Tok)
+	where
+		Tok: Token<E::T, Metadata=HostFnWeights<E::T>>,
+	{
+		self.gas_meter.adjust_gas(
+			charged_amount,
+			&self.ext.schedule().host_fn_weights,
+			adjusted_amount,
+		);
 	}
 
 	/// Read designated chunk from the sandbox memory.
@@ -774,11 +799,12 @@ define_env!(Env, <E: Ext>,
 			ctx.read_sandbox_memory_as(callee_ptr, callee_len)?;
 		let value: BalanceOf<<E as Ext>::T> = ctx.read_sandbox_memory_as(value_ptr, value_len)?;
 		let input_data = ctx.read_sandbox_memory(input_data_ptr, input_data_len)?;
-
 		if value > 0u32.into() {
 			ctx.charge_gas(RuntimeToken::CallSurchargeTransfer)?;
 		}
-
+		let charged = ctx.charge_gas(
+			RuntimeToken::CallSurchargeCodeSize(<E::T as Config>::MaxCodeSize::get())
+		)?;
 		let nested_gas_limit = if gas == 0 {
 			ctx.gas_meter.gas_left()
 		} else {
@@ -796,16 +822,20 @@ define_env!(Env, <E: Ext>,
 					)
 				}
 				// there is not enough gas to allocate for the nested call.
-				None => Err(Error::<<E as Ext>::T>::OutOfGas.into()),
+				None => Err((Error::<<E as Ext>::T>::OutOfGas.into(), 0)),
 			}
 		});
-
-		if let Ok(output) = &call_outcome {
+		let code_len = match &call_outcome {
+			Ok((_, len)) => len,
+			Err((_, len)) => len,
+		};
+		ctx.adjust_gas(charged, RuntimeToken::CallSurchargeCodeSize(*code_len));
+		if let Ok((output, _)) = &call_outcome {
 			ctx.write_sandbox_output(output_ptr, output_len_ptr, &output.data, true, |len| {
 				Some(RuntimeToken::CallCopyOut(len))
 			})?;
 		}
-		Ok(Runtime::<E>::exec_into_return_code(call_outcome)?)
+		Ok(Runtime::<E>::exec_into_return_code(call_outcome.map(|r| r.0).map_err(|r| r.0))?)
 	},
 
 	// Instantiate a contract with the specified code hash.
@@ -875,7 +905,9 @@ define_env!(Env, <E: Ext>,
 		let value: BalanceOf<<E as Ext>::T> = ctx.read_sandbox_memory_as(value_ptr, value_len)?;
 		let input_data = ctx.read_sandbox_memory(input_data_ptr, input_data_len)?;
 		let salt = ctx.read_sandbox_memory(salt_ptr, salt_len)?;
-
+		let charged = ctx.charge_gas(
+			RuntimeToken::InstantiateSurchargeCodeSize(<E::T as Config>::MaxCodeSize::get())
+		)?;
 		let nested_gas_limit = if gas == 0 {
 			ctx.gas_meter.gas_left()
 		} else {
@@ -894,10 +926,15 @@ define_env!(Env, <E: Ext>,
 					)
 				}
 				// there is not enough gas to allocate for the nested call.
-				None => Err(Error::<<E as Ext>::T>::OutOfGas.into()),
+				None => Err((Error::<<E as Ext>::T>::OutOfGas.into(), 0)),
 			}
 		});
-		if let Ok((address, output)) = &instantiate_outcome {
+		let code_len = match &instantiate_outcome {
+			Ok((_, _, code_len)) => code_len,
+			Err((_, code_len)) => code_len,
+		};
+		ctx.adjust_gas(charged, RuntimeToken::InstantiateSurchargeCodeSize(*code_len));
+		if let Ok((address, output, _)) = &instantiate_outcome {
 			if !output.flags.contains(ReturnFlags::REVERT) {
 				ctx.write_sandbox_output(
 					address_ptr, address_len_ptr, &address.encode(), true, already_charged,
@@ -907,7 +944,9 @@ define_env!(Env, <E: Ext>,
 				Some(RuntimeToken::InstantiateCopyOut(len))
 			})?;
 		}
-		Ok(Runtime::<E>::exec_into_return_code(instantiate_outcome.map(|(_id, retval)| retval))?)
+		Ok(Runtime::<E>::exec_into_return_code(
+			instantiate_outcome.map(|(_, retval, _)| retval).map_err(|(err, _)| err)
+		)?)
 	},
 
 	// Remove the calling account and transfer remaining balance.
@@ -935,7 +974,15 @@ define_env!(Env, <E: Ext>,
 		let beneficiary: <<E as Ext>::T as frame_system::Config>::AccountId =
 			ctx.read_sandbox_memory_as(beneficiary_ptr, beneficiary_len)?;
 
-		ctx.ext.terminate(&beneficiary)?;
+		let charged = ctx.charge_gas(
+			RuntimeToken::TerminateSurchargeCodeSize(<E::T as Config>::MaxCodeSize::get())
+		)?;
+		let (result, code_len) = match ctx.ext.terminate(&beneficiary) {
+			Ok(len) => (Ok(()), len),
+			Err((err, len)) => (Err(err), len),
+		};
+		ctx.adjust_gas(charged, RuntimeToken::TerminateSurchargeCodeSize(code_len));
+		result?;
 		Err(TrapReason::Termination)
 	},
 
@@ -1220,7 +1267,22 @@ define_env!(Env, <E: Ext>,
 			delta
 		};
 
-		ctx.ext.restore_to(dest, code_hash, rent_allowance, delta)?;
+		let max_len = <E::T as Config>::MaxCodeSize::get();
+		let charged = ctx.charge_gas(RuntimeToken::RestoreToSurchargeCodeSize {
+			caller_code: max_len,
+			tombstone_code: max_len,
+		})?;
+		let (result, caller_code, tombstone_code) = match ctx.ext.restore_to(
+			dest, code_hash, rent_allowance, delta
+		) {
+			Ok((code, tomb)) => (Ok(()), code, tomb),
+			Err((err, code, tomb)) => (Err(err), code, tomb),
+		};
+		ctx.adjust_gas(charged, RuntimeToken::RestoreToSurchargeCodeSize {
+			caller_code,
+			tombstone_code,
+		});
+		result?;
 		Err(TrapReason::Restoration)
 	},