Pre-Charge max size when contracts access storage (#10691)

* Fix seal_get_storage

* Fix seal_take_storage

* Add more benchmarks

* cargo run --quiet --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

* Fix seal_set_storage

* Fix seal_contains_storage and seal_clear_storage

* Fix benchmarks

* cargo run --quiet --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

* Get rid of mem::size_of in benchmarks

* Fix up code loading

* Apply suggestions from code review

Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>

* Fix test to call same function twice

* Replaced u32::MAX by SENTINEL const

* Fix seal_contains_storage benchmark

* cargo run --quiet --profile=production  --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

Co-authored-by: Parity Bot <admin@parity.io>
Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>

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

@@ -35,7 +35,7 @@ use crate::{
 	storage::Storage,
 	Pallet as Contracts, *,
 };
-use codec::Encode;
+use codec::{Encode, MaxEncodedLen};
 use frame_benchmarking::{account, benchmarks, whitelisted_caller};
 use frame_support::weights::Weight;
 use frame_system::RawOrigin;
@@ -778,9 +778,10 @@ benchmarks! {
 	seal_set_storage {
 		let r in 0 .. API_BENCHMARK_BATCHES;
 		let keys = (0 .. r * API_BENCHMARK_BATCH_SIZE)
-			.flat_map(|n| T::Hashing::hash_of(&n).as_ref().to_vec())
+			.map(|n| T::Hashing::hash_of(&n).as_ref().to_vec())
 			.collect::<Vec<_>>();
-		let key_len = sp_std::mem::size_of::<<T::Hashing as sp_runtime::traits::Hash>::Output>();
+		let key_len = keys.get(0).map(|i| i.len() as u32).unwrap_or(0);
+		let key_bytes = keys.iter().flatten().cloned().collect::<Vec<_>>();
 		let code = WasmModule::<T>::from(ModuleDefinition {
 			memory: Some(ImportedMemory::max::<T>()),
 			imported_functions: vec![ImportedFunction {
@@ -792,7 +793,7 @@ benchmarks! {
 			data_segments: vec![
 				DataSegment {
 					offset: 0,
-					value: keys,
+					value: key_bytes,
 				},
 			],
 			call_body: Some(body::repeated_dyn(r * API_BENCHMARK_BATCH_SIZE, vec![
@@ -805,13 +806,28 @@ benchmarks! {
 			.. Default::default()
 		});
 		let instance = Contract::<T>::new(code, vec![])?;
+		let info = instance.info()?;
+		for key in keys {
+			Storage::<T>::write(
+				&info.trie_id,
+				key.as_slice().try_into().map_err(|e| "Key has wrong length")?,
+				Some(vec![]),
+				None,
+				false,
+			)
+			.map_err(|_| "Failed to write to storage during setup.")?;
+		}
 		let origin = RawOrigin::Signed(instance.caller.clone());
 	}: call(origin, instance.addr, 0u32.into(), Weight::MAX, None, vec![])
 
-	seal_set_storage_per_kb {
+	#[skip_meta]
+	seal_set_storage_per_new_kb {
 		let n in 0 .. T::Schedule::get().limits.payload_len / 1024;
-		let key = T::Hashing::hash_of(&1u32).as_ref().to_vec();
+		let keys = (0 .. API_BENCHMARK_BATCH_SIZE)
-		let key_len = key.len();
+			.map(|n| T::Hashing::hash_of(&n).as_ref().to_vec())
+			.collect::<Vec<_>>();
+		let key_len = keys.get(0).map(|i| i.len() as u32).unwrap_or(0);
+		let key_bytes = keys.iter().flatten().cloned().collect::<Vec<_>>();
 		let code = WasmModule::<T>::from(ModuleDefinition {
 			memory: Some(ImportedMemory::max::<T>()),
 			imported_functions: vec![ImportedFunction {
@@ -823,19 +839,76 @@ benchmarks! {
 			data_segments: vec![
 				DataSegment {
 					offset: 0,
-					value: key,
+					value: key_bytes,
 				},
 			],
-			call_body: Some(body::repeated(API_BENCHMARK_BATCH_SIZE, &[
+			call_body: Some(body::repeated_dyn(API_BENCHMARK_BATCH_SIZE, vec![
-				Instruction::I32Const(0), // key_ptr
+				Counter(0, key_len as u32), // key_ptr
-				Instruction::I32Const(0), // value_ptr
+				Regular(Instruction::I32Const(0)), // value_ptr
-				Instruction::I32Const((n * 1024) as i32), // value_len
+				Regular(Instruction::I32Const((n * 1024) as i32)), // value_len
-				Instruction::Call(0),
+				Regular(Instruction::Call(0)),
-				Instruction::Drop,
+				Regular(Instruction::Drop),
 			])),
 			.. Default::default()
 		});
 		let instance = Contract::<T>::new(code, vec![])?;
+		let info = instance.info()?;
+		for key in keys {
+			Storage::<T>::write(
+				&info.trie_id,
+				key.as_slice().try_into().map_err(|e| "Key has wrong length")?,
+				Some(vec![]),
+				None,
+				false,
+			)
+			.map_err(|_| "Failed to write to storage during setup.")?;
+		}
+		let origin = RawOrigin::Signed(instance.caller.clone());
+	}: call(origin, instance.addr, 0u32.into(), Weight::MAX, None, vec![])
+
+	#[skip_meta]
+	seal_set_storage_per_old_kb {
+		let n in 0 .. T::Schedule::get().limits.payload_len / 1024;
+		let keys = (0 .. API_BENCHMARK_BATCH_SIZE)
+			.map(|n| T::Hashing::hash_of(&n).as_ref().to_vec())
+			.collect::<Vec<_>>();
+		let key_len = keys.get(0).map(|i| i.len() as u32).unwrap_or(0);
+		let key_bytes = keys.iter().flatten().cloned().collect::<Vec<_>>();
+		let code = WasmModule::<T>::from(ModuleDefinition {
+			memory: Some(ImportedMemory::max::<T>()),
+			imported_functions: vec![ImportedFunction {
+				module: "__unstable__",
+				name: "seal_set_storage",
+				params: vec![ValueType::I32, ValueType::I32, ValueType::I32],
+				return_type: Some(ValueType::I32),
+			}],
+			data_segments: vec![
+				DataSegment {
+					offset: 0,
+					value: key_bytes,
+				},
+			],
+			call_body: Some(body::repeated_dyn(API_BENCHMARK_BATCH_SIZE, vec![
+				Counter(0, key_len as u32), // key_ptr
+				Regular(Instruction::I32Const(0)), // value_ptr
+				Regular(Instruction::I32Const(0)), // value_len
+				Regular(Instruction::Call(0)),
+				Regular(Instruction::Drop),
+			])),
+			.. Default::default()
+		});
+		let instance = Contract::<T>::new(code, vec![])?;
+		let info = instance.info()?;
+		for key in keys {
+			Storage::<T>::write(
+				&info.trie_id,
+				key.as_slice().try_into().map_err(|e| "Key has wrong length")?,
+				Some(vec![42u8; (n * 1024) as usize]),
+				None,
+				false,
+			)
+			.map_err(|_| "Failed to write to storage during setup.")?;
+		}
 		let origin = RawOrigin::Signed(instance.caller.clone());
 	}: call(origin, instance.addr, 0u32.into(), Weight::MAX, None, vec![])
 
@@ -849,7 +922,7 @@ benchmarks! {
 			.map(|n| T::Hashing::hash_of(&n).as_ref().to_vec())
 			.collect::<Vec<_>>();
 		let key_bytes = keys.iter().flatten().cloned().collect::<Vec<_>>();
-		let key_len = sp_std::mem::size_of::<<T::Hashing as sp_runtime::traits::Hash>::Output>();
+		let key_len = keys.get(0).map(|i| i.len() as u32).unwrap_or(0);
 		let code = WasmModule::<T>::from(ModuleDefinition {
 			memory: Some(ImportedMemory::max::<T>()),
 			imported_functions: vec![ImportedFunction {
@@ -877,7 +950,7 @@ benchmarks! {
 			Storage::<T>::write(
 				&info.trie_id,
 				key.as_slice().try_into().map_err(|e| "Key has wrong length")?,
-				Some(vec![42; T::Schedule::get().limits.payload_len as usize]),
+				Some(vec![]),
 				None,
 				false,
 			)
@@ -887,6 +960,50 @@ benchmarks! {
 		let origin = RawOrigin::Signed(instance.caller.clone());
 	}: call(origin, instance.addr, 0u32.into(), Weight::MAX, None, vec![])
 
+	#[skip_meta]
+	seal_clear_storage_per_kb {
+		let n in 0 .. T::Schedule::get().limits.payload_len / 1024;
+		let keys = (0 .. API_BENCHMARK_BATCH_SIZE)
+			.map(|n| T::Hashing::hash_of(&n).as_ref().to_vec())
+			.collect::<Vec<_>>();
+		let key_len = keys.get(0).map(|i| i.len() as u32).unwrap_or(0);
+		let key_bytes = keys.iter().flatten().cloned().collect::<Vec<_>>();
+		let code = WasmModule::<T>::from(ModuleDefinition {
+			memory: Some(ImportedMemory::max::<T>()),
+			imported_functions: vec![ImportedFunction {
+				module: "__unstable__",
+				name: "seal_clear_storage",
+				params: vec![ValueType::I32],
+				return_type: Some(ValueType::I32),
+			}],
+			data_segments: vec![
+				DataSegment {
+					offset: 0,
+					value: key_bytes,
+				},
+			],
+			call_body: Some(body::repeated_dyn(API_BENCHMARK_BATCH_SIZE, vec![
+				Counter(0, key_len as u32), // key_ptr
+				Regular(Instruction::Call(0)),
+				Regular(Instruction::Drop),
+			])),
+			.. Default::default()
+		});
+		let instance = Contract::<T>::new(code, vec![])?;
+		let info = instance.info()?;
+		for key in keys {
+			Storage::<T>::write(
+				&info.trie_id,
+				key.as_slice().try_into().map_err(|e| "Key has wrong length")?,
+				Some(vec![42u8; (n * 1024) as usize]),
+				None,
+				false,
+			)
+			.map_err(|_| "Failed to write to storage during setup.")?;
+		}
+		let origin = RawOrigin::Signed(instance.caller.clone());
+	}: call(origin, instance.addr, 0u32.into(), Weight::MAX, None, vec![])
+
 	// We make sure that all storage accesses are to unique keys.
 	#[skip_meta]
 	seal_get_storage {
@@ -894,7 +1011,7 @@ benchmarks! {
 		let keys = (0 .. r * API_BENCHMARK_BATCH_SIZE)
 			.map(|n| T::Hashing::hash_of(&n).as_ref().to_vec())
 			.collect::<Vec<_>>();
-		let key_len = sp_std::mem::size_of::<<T::Hashing as sp_runtime::traits::Hash>::Output>();
+		let key_len = keys.get(0).map(|i| i.len() as u32).unwrap_or(0);
 		let key_bytes = keys.iter().flatten().cloned().collect::<Vec<_>>();
 		let key_bytes_len = key_bytes.len();
 		let code = WasmModule::<T>::from(ModuleDefinition {
@@ -940,6 +1057,58 @@ benchmarks! {
 		let origin = RawOrigin::Signed(instance.caller.clone());
 	}: call(origin, instance.addr, 0u32.into(), Weight::MAX, None, vec![])
 
+	#[skip_meta]
+	seal_get_storage_per_kb {
+		let n in 0 .. T::Schedule::get().limits.payload_len / 1024;
+		let keys = (0 .. API_BENCHMARK_BATCH_SIZE)
+			.map(|n| T::Hashing::hash_of(&n).as_ref().to_vec())
+			.collect::<Vec<_>>();
+		let key_len = keys.get(0).map(|i| i.len() as u32).unwrap_or(0);
+		let key_bytes = keys.iter().flatten().cloned().collect::<Vec<_>>();
+		let key_bytes_len = key_bytes.len();
+		let code = WasmModule::<T>::from(ModuleDefinition {
+			memory: Some(ImportedMemory::max::<T>()),
+			imported_functions: vec![ImportedFunction {
+				module: "seal0",
+				name: "seal_get_storage",
+				params: vec![ValueType::I32, ValueType::I32, ValueType::I32],
+				return_type: Some(ValueType::I32),
+			}],
+			data_segments: vec![
+				DataSegment {
+					offset: 0,
+					value: key_bytes,
+				},
+				DataSegment {
+					offset: key_bytes_len as u32,
+					value: T::Schedule::get().limits.payload_len.to_le_bytes().into(),
+				},
+			],
+			call_body: Some(body::repeated_dyn(API_BENCHMARK_BATCH_SIZE, vec![
+				Counter(0, key_len as u32), // key_ptr
+				Regular(Instruction::I32Const((key_bytes_len + 4) as i32)), // out_ptr
+				Regular(Instruction::I32Const(key_bytes_len as i32)), // out_len_ptr
+				Regular(Instruction::Call(0)),
+				Regular(Instruction::Drop),
+			])),
+			.. Default::default()
+		});
+		let instance = Contract::<T>::new(code, vec![])?;
+		let info = instance.info()?;
+		for key in keys {
+			Storage::<T>::write(
+				&info.trie_id,
+				key.as_slice().try_into().map_err(|e| "Key has wrong length")?,
+				Some(vec![42u8; (n * 1024) as usize]),
+				None,
+				false,
+			)
+			.map_err(|_| "Failed to write to storage during setup.")?;
+		}
+		<ContractInfoOf<T>>::insert(&instance.account_id, info.clone());
+		let origin = RawOrigin::Signed(instance.caller.clone());
+	}: call(origin, instance.addr, 0u32.into(), Weight::MAX, None, vec![])
+
 	// We make sure that all storage accesses are to unique keys.
 	#[skip_meta]
 	seal_contains_storage {
@@ -947,7 +1116,7 @@ benchmarks! {
 		let keys = (0 .. r * API_BENCHMARK_BATCH_SIZE)
 			.map(|n| T::Hashing::hash_of(&n).as_ref().to_vec())
 			.collect::<Vec<_>>();
-		let key_len = sp_std::mem::size_of::<<T::Hashing as sp_runtime::traits::Hash>::Output>();
+		let key_len = keys.get(0).map(|i| i.len() as u32).unwrap_or(0);
 		let key_bytes = keys.iter().flatten().cloned().collect::<Vec<_>>();
 		let key_bytes_len = key_bytes.len();
 		let code = WasmModule::<T>::from(ModuleDefinition {
@@ -977,7 +1146,7 @@ benchmarks! {
 			Storage::<T>::write(
 				&info.trie_id,
 				key.as_slice().try_into().map_err(|e| "Key has wrong length")?,
-				Some(vec![42; T::Schedule::get().limits.payload_len as usize]),
+				Some(vec![]),
 				None,
 				false,
 			)
@@ -987,48 +1156,47 @@ benchmarks! {
 		let origin = RawOrigin::Signed(instance.caller.clone());
 	}: call(origin, instance.addr, 0u32.into(), Weight::MAX, None, vec![])
 
-	seal_get_storage_per_kb {
+	#[skip_meta]
+	seal_contains_storage_per_kb {
 		let n in 0 .. T::Schedule::get().limits.payload_len / 1024;
-		let key = T::Hashing::hash_of(&1u32).as_ref().to_vec();
+		let keys = (0 .. API_BENCHMARK_BATCH_SIZE)
-		let key_len = key.len();
+			.map(|n| T::Hashing::hash_of(&n).as_ref().to_vec())
+			.collect::<Vec<_>>();
+		let key_len = keys.get(0).map(|i| i.len() as u32).unwrap_or(0);
+		let key_bytes = keys.iter().flatten().cloned().collect::<Vec<_>>();
 		let code = WasmModule::<T>::from(ModuleDefinition {
 			memory: Some(ImportedMemory::max::<T>()),
 			imported_functions: vec![ImportedFunction {
-				module: "seal0",
+				module: "__unstable__",
-				name: "seal_get_storage",
+				name: "seal_contains_storage",
-				params: vec![ValueType::I32, ValueType::I32, ValueType::I32],
+				params: vec![ValueType::I32],
 				return_type: Some(ValueType::I32),
 			}],
 			data_segments: vec![
 				DataSegment {
 					offset: 0,
-					value: key.clone(),
+					value: key_bytes,
-				},
-				DataSegment {
-					offset: key_len as u32,
-					value: T::Schedule::get().limits.payload_len.to_le_bytes().into(),
 				},
 			],
-			call_body: Some(body::repeated(API_BENCHMARK_BATCH_SIZE, &[
+			call_body: Some(body::repeated_dyn(API_BENCHMARK_BATCH_SIZE, vec![
-				// call at key_ptr
+				Counter(0, key_len as u32), // key_ptr
-				Instruction::I32Const(0), // key_ptr
+				Regular(Instruction::Call(0)),
-				Instruction::I32Const((key_len + 4) as i32), // out_ptr
+				Regular(Instruction::Drop),
-				Instruction::I32Const(key_len as i32), // out_len_ptr
-				Instruction::Call(0),
-				Instruction::Drop,
 			])),
 			.. Default::default()
 		});
 		let instance = Contract::<T>::new(code, vec![])?;
 		let info = instance.info()?;
-		Storage::<T>::write(
+		for key in keys {
-			&info.trie_id,
+			Storage::<T>::write(
-			key.as_slice().try_into().map_err(|e| "Key has wrong length")?,
+				&info.trie_id,
-			Some(vec![42u8; (n * 1024) as usize]),
+				key.as_slice().try_into().map_err(|e| "Key has wrong length")?,
-			None,
+				Some(vec![42u8; (n * 1024) as usize]),
-			false,
+				None,
-		)
+				false,
-		.map_err(|_| "Failed to write to storage during setup.")?;
+			)
+			.map_err(|_| "Failed to write to storage during setup.")?;
+		}
 		<ContractInfoOf<T>>::insert(&instance.account_id, info.clone());
 		let origin = RawOrigin::Signed(instance.caller.clone());
 	}: call(origin, instance.addr, 0u32.into(), Weight::MAX, None, vec![])
@@ -1039,7 +1207,7 @@ benchmarks! {
 		let keys = (0 .. r * API_BENCHMARK_BATCH_SIZE)
 			.map(|n| T::Hashing::hash_of(&n).as_ref().to_vec())
 			.collect::<Vec<_>>();
-		let key_len = sp_std::mem::size_of::<<T::Hashing as sp_runtime::traits::Hash>::Output>();
+		let key_len = keys.get(0).map(|i| i.len() as u32).unwrap_or(0);
 		let key_bytes = keys.iter().flatten().cloned().collect::<Vec<_>>();
 		let key_bytes_len = key_bytes.len();
 		let code = WasmModule::<T>::from(ModuleDefinition {
@@ -1091,7 +1259,7 @@ benchmarks! {
 		let keys = (0 .. API_BENCHMARK_BATCH_SIZE)
 			.map(|n| T::Hashing::hash_of(&n).as_ref().to_vec())
 			.collect::<Vec<_>>();
-		let key_len = sp_std::mem::size_of::<<T::Hashing as sp_runtime::traits::Hash>::Output>();
+		let key_len = keys.get(0).map(|i| i.len() as u32).unwrap_or(0);
 		let key_bytes = keys.iter().flatten().cloned().collect::<Vec<_>>();
 		let key_bytes_len = key_bytes.len();
 		let code = WasmModule::<T>::from(ModuleDefinition {
@@ -1238,7 +1406,7 @@ benchmarks! {
 				Regular(Instruction::I32Const(value_len as i32)), // value_len
 				Regular(Instruction::I32Const(0)), // input_data_ptr
 				Regular(Instruction::I32Const(0)), // input_data_len
-				Regular(Instruction::I32Const(u32::max_value() as i32)), // output_ptr
+				Regular(Instruction::I32Const(SENTINEL as i32)), // output_ptr
 				Regular(Instruction::I32Const(0)), // output_len_ptr
 				Regular(Instruction::Call(0)),
 				Regular(Instruction::Drop),
@@ -1361,7 +1529,7 @@ benchmarks! {
 		assert!(value > 0u32.into());
 		let value_bytes = value.encode();
 		let value_len = value_bytes.len();
-		let addr_len = sp_std::mem::size_of::<T::AccountId>();
+		let addr_len = T::AccountId::max_encoded_len();
 
 		// offsets where to place static data in contract memory
 		let value_offset = 0;
@@ -1415,7 +1583,7 @@ benchmarks! {
 				Regular(Instruction::I32Const(0)), // input_data_len
 				Regular(Instruction::I32Const(addr_offset as i32)), // address_ptr
 				Regular(Instruction::I32Const(addr_len_offset as i32)), // address_len_ptr
-				Regular(Instruction::I32Const(u32::max_value() as i32)), // output_ptr
+				Regular(Instruction::I32Const(SENTINEL as i32)), // output_ptr
 				Regular(Instruction::I32Const(0)), // output_len_ptr
 				Regular(Instruction::I32Const(0)), // salt_ptr
 				Regular(Instruction::I32Const(0)), // salt_ptr_len
@@ -1485,7 +1653,7 @@ benchmarks! {
 		assert!(value > 0u32.into());
 		let value_bytes = value.encode();
 		let value_len = value_bytes.len();
-		let addr_len = sp_std::mem::size_of::<T::AccountId>();
+		let addr_len = T::AccountId::max_encoded_len();
 
 		// offsets where to place static data in contract memory
 		let input_offset = 0;

substrate/frame/contracts/src/chain_extension.rs

@@ -329,7 +329,7 @@ where
 	///
 	/// If the contract supplied buffer is smaller than the passed `buffer` an `Err` is returned.
 	/// If `allow_skip` is set to true the contract is allowed to skip the copying of the buffer
-	/// by supplying the guard value of `u32::MAX` as `out_ptr`. The
+	/// by supplying the guard value of `pallet-contracts::SENTINEL` as `out_ptr`. The
 	/// `weight_per_byte` is only charged when the write actually happens and is not skipped or
 	/// failed due to a too small output buffer.
 	pub fn write(

substrate/frame/contracts/src/exec.rs

@@ -24,7 +24,7 @@ use crate::{
 use frame_support::{
 	dispatch::{DispatchError, DispatchResult, DispatchResultWithPostInfo, Dispatchable},
 	storage::{with_transaction, TransactionOutcome},
-	traits::{Contains, Currency, ExistenceRequirement, Get, OriginTrait, Randomness, Time},
+	traits::{Contains, Currency, ExistenceRequirement, OriginTrait, Randomness, Time},
 	weights::Weight,
 };
 use frame_system::RawOrigin;
@@ -140,11 +140,11 @@ pub trait Ext: sealing::Sealed {
 	/// was deleted.
 	fn get_storage(&mut self, key: &StorageKey) -> Option<Vec<u8>>;
 
-	/// Returns true iff some storage entry exists under the supplied `key`
+	/// Returns `Some(len)` (in bytes) if a storage item exists at `key`.
 	///
-	/// Returns `false` if the `key` wasn't previously set by `set_storage` or
+	/// Returns `None` if the `key` wasn't previously set by `set_storage` or
 	/// was deleted.
-	fn contains_storage(&mut self, key: &StorageKey) -> bool;
+	fn get_storage_size(&mut self, key: &StorageKey) -> Option<u32>;
 
 	/// Sets the storage entry by the given key to the specified value. If `value` is `None` then
 	/// the storage entry is deleted.
@@ -996,8 +996,8 @@ where
 		Storage::<T>::read(&self.top_frame_mut().contract_info().trie_id, key)
 	}
 
-	fn contains_storage(&mut self, key: &StorageKey) -> bool {
+	fn get_storage_size(&mut self, key: &StorageKey) -> Option<u32> {
-		Storage::<T>::contains(&self.top_frame_mut().contract_info().trie_id, key)
+		Storage::<T>::size(&self.top_frame_mut().contract_info().trie_id, key)
 	}
 
 	fn set_storage(
@@ -1056,7 +1056,7 @@ where
 	}
 
 	fn max_value_size(&self) -> u32 {
-		T::Schedule::get().limits.payload_len
+		self.schedule.limits.payload_len
 	}
 
 	fn get_weight_price(&self, weight: Weight) -> BalanceOf<Self::T> {
@@ -2432,16 +2432,16 @@ mod tests {
 	}
 
 	#[test]
-	fn contains_storage_works() {
+	fn get_storage_size_works() {
 		let code_hash = MockLoader::insert(Call, |ctx, _| {
 			assert_eq!(
 				ctx.ext.set_storage([1; 32], Some(vec![1, 2, 3]), false),
 				Ok(WriteOutcome::New)
 			);
 			assert_eq!(ctx.ext.set_storage([2; 32], Some(vec![]), false), Ok(WriteOutcome::New));
-			assert_eq!(ctx.ext.contains_storage(&[1; 32]), true);
+			assert_eq!(ctx.ext.get_storage_size(&[1; 32]), Some(3));
-			assert_eq!(ctx.ext.contains_storage(&[1; 32]), true);
+			assert_eq!(ctx.ext.get_storage_size(&[2; 32]), Some(0));
-			assert_eq!(ctx.ext.contains_storage(&[3; 32]), false);
+			assert_eq!(ctx.ext.get_storage_size(&[3; 32]), None);
 
 			exec_success()
 		});

substrate/frame/contracts/src/lib.rs

@@ -136,6 +136,14 @@ type BalanceOf<T> =
 /// The current storage version.
 const STORAGE_VERSION: StorageVersion = StorageVersion::new(6);
 
+/// Used as a sentinel value when reading and writing contract memory.
+///
+/// It is usually used to signal `None` to a contract when only a primitive is allowed
+/// and we don't want to go through encoding a full Rust type. Using `u32::Max` is a safe
+/// sentinel because contracts are never allowed to use such a large amount of resources
+/// that this value makes sense for a memory location or length.
+const SENTINEL: u32 = u32::MAX;
+
 /// Provides the contract address generation method.
 ///
 /// See [`DefaultAddressGenerator`] for the default implementation.
@@ -831,7 +839,7 @@ where
 		module: &mut PrefabWasmModule<T>,
 		schedule: &Schedule<T>,
 	) -> frame_support::dispatch::DispatchResult {
-		self::wasm::reinstrument(module, schedule)
+		self::wasm::reinstrument(module, schedule).map(|_| ())
 	}
 
 	/// Internal function that does the actual call.

substrate/frame/contracts/src/schedule.rs

@@ -316,15 +316,24 @@ pub struct HostFnWeights<T: Config> {
 	/// Weight of calling `seal_set_storage`.
 	pub set_storage: Weight,
 
-	/// Weight per byte of an item stored with `seal_set_storage`.
+	/// Weight per written byten of an item stored with `seal_set_storage`.
-	pub set_storage_per_byte: Weight,
+	pub set_storage_per_new_byte: Weight,
+
+	/// Weight per overwritten byte of an item stored with `seal_set_storage`.
+	pub set_storage_per_old_byte: Weight,
 
 	/// Weight of calling `seal_clear_storage`.
 	pub clear_storage: Weight,
 
+	/// Weight of calling `seal_clear_storage` per byte of the stored item.
+	pub clear_storage_per_byte: Weight,
+
 	/// Weight of calling `seal_contains_storage`.
 	pub contains_storage: Weight,
 
+	/// Weight of calling `seal_contains_storage` per byte of the stored item.
+	pub contains_storage_per_byte: Weight,
+
 	/// Weight of calling `seal_get_storage`.
 	pub get_storage: Weight,
 
@@ -586,9 +595,12 @@ impl<T: Config> Default for HostFnWeights<T> {
 			),
 			debug_message: cost_batched!(seal_debug_message),
 			set_storage: cost_batched!(seal_set_storage),
-			set_storage_per_byte: cost_byte_batched!(seal_set_storage_per_kb),
+			set_storage_per_new_byte: cost_byte_batched!(seal_set_storage_per_new_kb),
+			set_storage_per_old_byte: cost_byte_batched!(seal_set_storage_per_old_kb),
 			clear_storage: cost_batched!(seal_clear_storage),
+			clear_storage_per_byte: cost_byte_batched!(seal_clear_storage_per_kb),
 			contains_storage: cost_batched!(seal_contains_storage),
+			contains_storage_per_byte: cost_byte_batched!(seal_contains_storage_per_kb),
 			get_storage: cost_batched!(seal_get_storage),
 			get_storage_per_byte: cost_byte_batched!(seal_get_storage_per_kb),
 			take_storage: cost_batched!(seal_take_storage),

substrate/frame/contracts/src/storage.rs

@@ -22,7 +22,7 @@ pub mod meter;
 use crate::{
 	exec::{AccountIdOf, StorageKey},
 	weights::WeightInfo,
-	BalanceOf, CodeHash, Config, ContractInfoOf, DeletionQueue, Error, TrieId,
+	BalanceOf, CodeHash, Config, ContractInfoOf, DeletionQueue, Error, TrieId, SENTINEL,
 };
 use codec::{Decode, Encode};
 use frame_support::{
@@ -87,6 +87,33 @@ pub enum WriteOutcome {
 	Taken(Vec<u8>),
 }
 
+impl WriteOutcome {
+	/// Extracts the size of the overwritten value or `0` if there
+	/// was no value in storage.
+	pub fn old_len(&self) -> u32 {
+		match self {
+			Self::New => 0,
+			Self::Overwritten(len) => *len,
+			Self::Taken(value) => value.len() as u32,
+		}
+	}
+
+	/// Extracts the size of the overwritten value or `SENTINEL` if there
+	/// was no value in storage.
+	///
+	/// # Note
+	///
+	/// We cannot use `0` as sentinel value because there could be a zero sized
+	/// storage entry which is different from a non existing one.
+	pub fn old_len_with_sentinel(&self) -> u32 {
+		match self {
+			Self::New => SENTINEL,
+			Self::Overwritten(len) => *len,
+			Self::Taken(value) => value.len() as u32,
+		}
+	}
+}
+
 pub struct Storage<T>(PhantomData<T>);
 
 impl<T> Storage<T>
@@ -102,9 +129,12 @@ where
 		child::get_raw(&child_trie_info(trie_id), &blake2_256(key))
 	}
 
-	/// Returns `true` iff the `key` exists in storage.
+	/// Returns `Some(len)` (in bytes) if a storage item exists at `key`.
-	pub fn contains(trie_id: &TrieId, key: &StorageKey) -> bool {
+	///
-		child::exists(&child_trie_info(trie_id), &blake2_256(key))
+	/// Returns `None` if the `key` wasn't previously set by `set_storage` or
+	/// was deleted.
+	pub fn size(trie_id: &TrieId, key: &StorageKey) -> Option<u32> {
+		child::len(&child_trie_info(trie_id), &blake2_256(key))
 	}
 
 	/// Update a storage entry into a contract's kv storage.

substrate/frame/contracts/src/tests.rs

@@ -294,7 +294,7 @@ pub const BOB: AccountId32 = AccountId32::new([2u8; 32]);
 pub const CHARLIE: AccountId32 = AccountId32::new([3u8; 32]);
 pub const DJANGO: AccountId32 = AccountId32::new([4u8; 32]);
 
-pub const GAS_LIMIT: Weight = 10_000_000_000;
+pub const GAS_LIMIT: Weight = 100_000_000_000;
 
 pub struct ExtBuilder {
 	existential_deposit: u64,

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

@@ -38,7 +38,6 @@ use crate::{
 use frame_support::{
 	dispatch::{DispatchError, DispatchResult},
 	ensure,
-	storage::StorageMap,
 	traits::ReservableCurrency,
 };
 use sp_core::crypto::UncheckedFrom;
@@ -149,52 +148,38 @@ pub fn load<T: Config>(
 where
 	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>,
 {
-	gas_meter.charge(CodeToken::Load(estimate_code_size::<T, CodeStorage<T>, _>(&code_hash)?))?;
+	let charged = gas_meter.charge(CodeToken::Load(schedule.limits.code_len))?;
 
 	let mut prefab_module =
 		<CodeStorage<T>>::get(code_hash).ok_or_else(|| Error::<T>::CodeNotFound)?;
+	gas_meter.adjust_gas(charged, CodeToken::Load(prefab_module.code.len() as u32));
 	prefab_module.code_hash = code_hash;
 
 	if prefab_module.instruction_weights_version < schedule.instruction_weights.version {
 		// The instruction weights have changed.
 		// We need to re-instrument the code with the new instruction weights.
-		gas_meter.charge(CodeToken::Reinstrument(estimate_code_size::<T, PristineCode<T>, _>(
+		let charged = gas_meter.charge(CodeToken::Reinstrument(schedule.limits.code_len))?;
-			&code_hash,
+		let code_size = reinstrument(&mut prefab_module, schedule)?;
-		)?))?;
+		gas_meter.adjust_gas(charged, CodeToken::Reinstrument(code_size));
-		reinstrument(&mut prefab_module, schedule)?;
 	}
 
 	Ok(prefab_module)
 }
 
 /// Instruments the passed prefab wasm module with the supplied schedule.
+///
+/// Returns the size in bytes of the uninstrumented code.
 pub fn reinstrument<T: Config>(
 	prefab_module: &mut PrefabWasmModule<T>,
 	schedule: &Schedule<T>,
-) -> Result<(), DispatchError> {
+) -> Result<u32, DispatchError> {
 	let original_code =
 		<PristineCode<T>>::get(&prefab_module.code_hash).ok_or_else(|| Error::<T>::CodeNotFound)?;
+	let original_code_len = original_code.len();
 	prefab_module.code = prepare::reinstrument_contract::<T>(original_code, schedule)?;
 	prefab_module.instruction_weights_version = schedule.instruction_weights.version;
 	<CodeStorage<T>>::insert(&prefab_module.code_hash, &*prefab_module);
-	Ok(())
+	Ok(original_code_len as u32)
-}
-
-/// Get the size of the code stored at `code_hash` without loading it.
-///
-/// The returned value is slightly too large when using it for the [`PrefabWasmModule`]
-/// because it has other fields in addition to the code itself. However, those are negligible
-/// when compared to the code size. Additionally, charging too much weight is completely safe.
-fn estimate_code_size<T, M, V>(code_hash: &CodeHash<T>) -> Result<u32, DispatchError>
-where
-	T: Config,
-	M: StorageMap<CodeHash<T>, V>,
-	V: codec::FullCodec,
-{
-	let key = M::hashed_key_for(code_hash);
-	let mut data = [0u8; 0];
-	let len = sp_io::storage::read(&key, &mut data, 0).ok_or_else(|| Error::<T>::CodeNotFound)?;
-	Ok(len)
 }
 
 /// Costs for operations that are related to code handling.

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

@@ -385,8 +385,8 @@ mod tests {
 		fn get_storage(&mut self, key: &StorageKey) -> Option<Vec<u8>> {
 			self.storage.get(key).cloned()
 		}
-		fn contains_storage(&mut self, key: &StorageKey) -> bool {
+		fn get_storage_size(&mut self, key: &StorageKey) -> Option<u32> {
-			self.storage.contains_key(key)
+			self.storage.get(key).map(|val| val.len() as u32)
 		}
 		fn set_storage(
 			&mut self,
@@ -2023,7 +2023,7 @@ mod tests {
 		// value did not exist before -> sentinel returned
 		let input = ([1u8; 32], [42u8, 48]).encode();
 		let result = execute(CODE, input, &mut ext).unwrap();
-		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), u32::MAX);
+		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), crate::SENTINEL);
 		assert_eq!(ext.storage.get(&[1u8; 32]).unwrap(), &[42u8, 48]);
 
 		// value do exist -> length of old value returned
@@ -2083,7 +2083,7 @@ mod tests {
 
 		// value does not exist -> sentinel returned
 		let result = execute(CODE, [3u8; 32].encode(), &mut ext).unwrap();
-		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), u32::MAX);
+		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), crate::SENTINEL);
 		assert_eq!(ext.storage.get(&[3u8; 32]), None);
 
 		// value did exist -> length returned
@@ -2228,25 +2228,16 @@ mod tests {
 		ext.storage.insert([1u8; 32], vec![42u8]);
 		ext.storage.insert([2u8; 32], vec![]);
 
-		// value does not exist -> error returned
+		// value does not exist -> sentinel value returned
 		let result = execute(CODE, [3u8; 32].encode(), &mut ext).unwrap();
-		assert_eq!(
+		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), crate::SENTINEL);
-			u32::from_le_bytes(result.data.0.try_into().unwrap()),
-			ReturnCode::KeyNotFound as u32
-		);
 
 		// value did exist -> success
 		let result = execute(CODE, [1u8; 32].encode(), &mut ext).unwrap();
-		assert_eq!(
+		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), 1,);
-			u32::from_le_bytes(result.data.0.try_into().unwrap()),
-			ReturnCode::Success as u32
-		);
 
 		// value did exist -> success (zero sized type)
 		let result = execute(CODE, [2u8; 32].encode(), &mut ext).unwrap();
-		assert_eq!(
+		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), 0,);
-			u32::from_le_bytes(result.data.0.try_into().unwrap()),
-			ReturnCode::Success as u32
-		);
 	}
 }

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

@@ -21,9 +21,8 @@ use crate::{
 	exec::{ExecError, ExecResult, Ext, StorageKey, TopicOf},
 	gas::{ChargedAmount, Token},
 	schedule::HostFnWeights,
-	storage::WriteOutcome,
 	wasm::env_def::ConvertibleToWasm,
-	BalanceOf, CodeHash, Config, Error,
+	BalanceOf, CodeHash, Config, Error, SENTINEL,
 };
 use bitflags::bitflags;
 use codec::{Decode, DecodeAll, Encode, MaxEncodedLen};
@@ -169,23 +168,18 @@ pub enum RuntimeCosts {
 	DepositEvent { num_topic: u32, len: u32 },
 	/// Weight of calling `seal_debug_message`.
 	DebugMessage,
-	/// Weight of calling `seal_set_storage` for the given storage item size.
+	/// Weight of calling `seal_set_storage` for the given storage item sizes.
-	SetStorage(u32),
+	SetStorage { old_bytes: u32, new_bytes: u32 },
-	/// Weight of calling `seal_clear_storage`.
+	/// Weight of calling `seal_clear_storage` per cleared byte.
-	ClearStorage,
+	ClearStorage(u32),
-	/// Weight of calling `seal_contains_storage`.
+	/// Weight of calling `seal_contains_storage` per byte of the checked item.
 	#[cfg(feature = "unstable-interface")]
-	ContainsStorage,
+	ContainsStorage(u32),
-	/// Weight of calling `seal_get_storage` without output weight.
+	/// Weight of calling `seal_get_storage` with the specified size in storage.
-	GetStorageBase,
+	GetStorage(u32),
-	/// Weight of an item received via `seal_get_storage` for the given size.
+	/// Weight of calling `seal_take_storage` for the given size.
-	GetStorageCopyOut(u32),
-	/// Weight of calling `seal_take_storage` without output weight.
 	#[cfg(feature = "unstable-interface")]
-	TakeStorageBase,
+	TakeStorage(u32),
-	/// Weight of an item received via `seal_take_storage` for the given size.
-	#[cfg(feature = "unstable-interface")]
-	TakeStorageCopyOut(u32),
 	/// Weight of calling `seal_transfer`.
 	Transfer,
 	/// Weight of calling `seal_call` for the given input size.
@@ -249,17 +243,23 @@ impl RuntimeCosts {
 				.saturating_add(s.deposit_event_per_topic.saturating_mul(num_topic.into()))
 				.saturating_add(s.deposit_event_per_byte.saturating_mul(len.into())),
 			DebugMessage => s.debug_message,
-			SetStorage(len) =>
+			SetStorage { new_bytes, old_bytes } => s
-				s.set_storage.saturating_add(s.set_storage_per_byte.saturating_mul(len.into())),
+				.set_storage
-			ClearStorage => s.clear_storage,
+				.saturating_add(s.set_storage_per_new_byte.saturating_mul(new_bytes.into()))
+				.saturating_add(s.set_storage_per_old_byte.saturating_mul(old_bytes.into())),
+			ClearStorage(len) => s
+				.clear_storage
+				.saturating_add(s.clear_storage_per_byte.saturating_mul(len.into())),
 			#[cfg(feature = "unstable-interface")]
-			ContainsStorage => s.contains_storage,
+			ContainsStorage(len) => s
-			GetStorageBase => s.get_storage,
+				.contains_storage
-			GetStorageCopyOut(len) => s.get_storage_per_byte.saturating_mul(len.into()),
+				.saturating_add(s.contains_storage_per_byte.saturating_mul(len.into())),
+			GetStorage(len) =>
+				s.get_storage.saturating_add(s.get_storage_per_byte.saturating_mul(len.into())),
 			#[cfg(feature = "unstable-interface")]
-			TakeStorageBase => s.take_storage,
+			TakeStorage(len) => s
-			#[cfg(feature = "unstable-interface")]
+				.take_storage
-			TakeStorageCopyOut(len) => s.take_storage_per_byte.saturating_mul(len.into()),
+				.saturating_add(s.take_storage_per_byte.saturating_mul(len.into())),
 			Transfer => s.transfer,
 			CallBase(len) =>
 				s.call.saturating_add(s.call_per_input_byte.saturating_mul(len.into())),
@@ -534,7 +534,7 @@ where
 	/// length of the buffer located at `out_ptr`. If that buffer is large enough the actual
 	/// `buf.len()` is written to this location.
 	///
-	/// If `out_ptr` is set to the sentinel value of `u32::MAX` and `allow_skip` is true the
+	/// If `out_ptr` is set to the sentinel value of `SENTINEL` and `allow_skip` is true the
 	/// operation is skipped and `Ok` is returned. This is supposed to help callers to make copying
 	/// output optional. For example to skip copying back the output buffer of an `seal_call`
 	/// when the caller is not interested in the result.
@@ -553,7 +553,7 @@ where
 		allow_skip: bool,
 		create_token: impl FnOnce(u32) -> Option<RuntimeCosts>,
 	) -> Result<(), DispatchError> {
-		if allow_skip && out_ptr == u32::MAX {
+		if allow_skip && out_ptr == SENTINEL {
 			return Ok(())
 		}
 
@@ -648,48 +648,36 @@ where
 		}
 	}
 
-	/// Extracts the size of the overwritten value or `u32::MAX` if there
-	/// was no value in storage.
-	///
-	/// # Note
-	///
-	/// We cannot use `0` as sentinel value because there could be a zero sized
-	/// storage entry which is different from a non existing one.
-	fn overwritten_len(outcome: WriteOutcome) -> u32 {
-		match outcome {
-			WriteOutcome::New => u32::MAX,
-			WriteOutcome::Overwritten(len) => len,
-			WriteOutcome::Taken(value) => value.len() as u32,
-		}
-	}
-
 	fn set_storage(
 		&mut self,
 		key_ptr: u32,
 		value_ptr: u32,
 		value_len: u32,
 	) -> Result<u32, TrapReason> {
-		self.charge_gas(RuntimeCosts::SetStorage(value_len))?;
+		let max_size = self.ext.max_value_size();
-		if value_len > self.ext.max_value_size() {
+		let charged = self
+			.charge_gas(RuntimeCosts::SetStorage { new_bytes: value_len, old_bytes: max_size })?;
+		if value_len > max_size {
 			Err(Error::<E::T>::ValueTooLarge)?;
 		}
 		let mut key: StorageKey = [0; 32];
 		self.read_sandbox_memory_into_buf(key_ptr, &mut key)?;
 		let value = Some(self.read_sandbox_memory(value_ptr, value_len)?);
-		self.ext
+		let write_outcome = self.ext.set_storage(key, value, false)?;
-			.set_storage(key, value, false)
+		self.adjust_gas(
-			.map(Self::overwritten_len)
+			charged,
-			.map_err(Into::into)
+			RuntimeCosts::SetStorage { new_bytes: value_len, old_bytes: write_outcome.old_len() },
+		);
+		Ok(write_outcome.old_len_with_sentinel())
 	}
 
 	fn clear_storage(&mut self, key_ptr: u32) -> Result<u32, TrapReason> {
-		self.charge_gas(RuntimeCosts::ClearStorage)?;
+		let charged = self.charge_gas(RuntimeCosts::ClearStorage(self.ext.max_value_size()))?;
 		let mut key: StorageKey = [0; 32];
 		self.read_sandbox_memory_into_buf(key_ptr, &mut key)?;
-		self.ext
+		let outcome = self.ext.set_storage(key, None, false)?;
-			.set_storage(key, None, false)
+		self.adjust_gas(charged, RuntimeCosts::ClearStorage(outcome.old_len()));
-			.map(Self::overwritten_len)
+		Ok(outcome.old_len_with_sentinel())
-			.map_err(Into::into)
 	}
 
 	fn call(
@@ -827,7 +815,7 @@ define_env!(Env, <E: Ext>,
 	// # Return Value
 	//
 	// Returns the size of the pre-existing value at the specified key if any. Otherwise
-	// `u32::MAX` is returned as a sentinel value.
+	// `SENTINEL` is returned as a sentinel value.
 	[__unstable__] seal_set_storage(ctx, key_ptr: u32, value_ptr: u32, value_len: u32) -> u32 => {
 		ctx.set_storage(key_ptr, value_ptr, value_len)
 	},
@@ -849,7 +837,7 @@ define_env!(Env, <E: Ext>,
 	// # Return Value
 	//
 	// Returns the size of the pre-existing value at the specified key if any. Otherwise
-	// `u32::MAX` is returned as a sentinel value.
+	// `SENTINEL` is returned as a sentinel value.
 	[__unstable__] seal_clear_storage(ctx, key_ptr: u32) -> u32 => {
 		ctx.clear_storage(key_ptr).map_err(Into::into)
 	},
@@ -867,39 +855,39 @@ define_env!(Env, <E: Ext>,
 	//
 	// `ReturnCode::KeyNotFound`
 	[seal0] seal_get_storage(ctx, key_ptr: u32, out_ptr: u32, out_len_ptr: u32) -> ReturnCode => {
-		ctx.charge_gas(RuntimeCosts::GetStorageBase)?;
+		let charged = ctx.charge_gas(RuntimeCosts::GetStorage(ctx.ext.max_value_size()))?;
 		let mut key: StorageKey = [0; 32];
 		ctx.read_sandbox_memory_into_buf(key_ptr, &mut key)?;
 		if let Some(value) = ctx.ext.get_storage(&key) {
-			ctx.write_sandbox_output(out_ptr, out_len_ptr, &value, false, |len| {
+			ctx.adjust_gas(charged, RuntimeCosts::GetStorage(value.len() as u32));
-				Some(RuntimeCosts::GetStorageCopyOut(len))
+			ctx.write_sandbox_output(out_ptr, out_len_ptr, &value, false, already_charged)?;
-			})?;
 			Ok(ReturnCode::Success)
 		} else {
+			ctx.adjust_gas(charged, RuntimeCosts::GetStorage(0));
 			Ok(ReturnCode::KeyNotFound)
 		}
 	},
 
 	// Checks whether there is a value stored under the given key.
 	//
-	// Returns `ReturnCode::Success` if there is a key in storage. Otherwise an error
-	// is returned.
-	//
 	// # Parameters
 	//
 	// - `key_ptr`: pointer into the linear memory where the key of the requested value is placed.
 	//
-	// # Errors
+	// # Return Value
 	//
-	// `ReturnCode::KeyNotFound`
+	// Returns the size of the pre-existing value at the specified key if any. Otherwise
-	[__unstable__] seal_contains_storage(ctx, key_ptr: u32) -> ReturnCode => {
+	// `SENTINEL` is returned as a sentinel value.
-		ctx.charge_gas(RuntimeCosts::ContainsStorage)?;
+	[__unstable__] seal_contains_storage(ctx, key_ptr: u32) -> u32 => {
+		let charged = ctx.charge_gas(RuntimeCosts::ContainsStorage(ctx.ext.max_value_size()))?;
 		let mut key: StorageKey = [0; 32];
 		ctx.read_sandbox_memory_into_buf(key_ptr, &mut key)?;
-		if ctx.ext.contains_storage(&key) {
+		if let Some(len) = ctx.ext.get_storage_size(&key) {
-			Ok(ReturnCode::Success)
+			ctx.adjust_gas(charged, RuntimeCosts::ContainsStorage(len));
+			Ok(len)
 		} else {
-			Ok(ReturnCode::KeyNotFound)
+			ctx.adjust_gas(charged, RuntimeCosts::ContainsStorage(0));
+			Ok(SENTINEL)
 		}
 	},
 
@@ -916,15 +904,15 @@ define_env!(Env, <E: Ext>,
 	//
 	// `ReturnCode::KeyNotFound`
 	[__unstable__] seal_take_storage(ctx, key_ptr: u32, out_ptr: u32, out_len_ptr: u32) -> ReturnCode => {
-		ctx.charge_gas(RuntimeCosts::TakeStorageBase)?;
+		let charged = ctx.charge_gas(RuntimeCosts::TakeStorage(ctx.ext.max_value_size()))?;
 		let mut key: StorageKey = [0; 32];
 		ctx.read_sandbox_memory_into_buf(key_ptr, &mut key)?;
-		if let WriteOutcome::Taken(value) = ctx.ext.set_storage(key, None, true)? {
+		if let crate::storage::WriteOutcome::Taken(value) = ctx.ext.set_storage(key, None, true)? {
-			ctx.write_sandbox_output(out_ptr, out_len_ptr, &value, false, |len| {
+			ctx.adjust_gas(charged, RuntimeCosts::TakeStorage(value.len() as u32));
-				Some(RuntimeCosts::TakeStorageCopyOut(len))
+			ctx.write_sandbox_output(out_ptr, out_len_ptr, &value, false, already_charged)?;
-			})?;
 			Ok(ReturnCode::Success)
 		} else {
+			ctx.adjust_gas(charged, RuntimeCosts::TakeStorage(0));
 			Ok(ReturnCode::KeyNotFound)
 		}
 	},
@@ -1006,7 +994,7 @@ define_env!(Env, <E: Ext>,
 	//
 	// The callees output buffer is copied to `output_ptr` and its length to `output_len_ptr`.
 	// The copy of the output buffer can be skipped by supplying the sentinel value
-	// of `u32::MAX` to `output_ptr`.
+	// of `SENTINEL` to `output_ptr`.
 	//
 	// # Parameters
 	//
@@ -1103,7 +1091,7 @@ define_env!(Env, <E: Ext>,
 	// by the code hash. The address of this new account is copied to `address_ptr` and its length
 	// to `address_len_ptr`. The constructors output buffer is copied to `output_ptr` and its
 	// length to `output_len_ptr`. The copy of the output buffer and address can be skipped by
-	// supplying the sentinel value of `u32::MAX` to `output_ptr` or `address_ptr`.
+	// supplying the sentinel value of `SENTINEL` to `output_ptr` or `address_ptr`.
 	//
 	// `value` must be at least the minimum balance. Otherwise the instantiation fails and the
 	// contract is not created.