Add additionally functionality to contracts storage interface (#10497)

* Add new versions for storage access host functions

* Improve docs

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

@@ -134,7 +134,7 @@ where
 	fn store(&self, items: &Vec<(StorageKey, Vec<u8>)>) -> Result<(), &'static str> {
 		let info = self.info()?;
 		for item in items {
-			Storage::<T>::write(&info.trie_id, &item.0, Some(item.1.clone()), None)
+			Storage::<T>::write(&info.trie_id, &item.0, Some(item.1.clone()), None, false)
 				.map_err(|_| "Failed to write storage to restoration dest")?;
 		}
 		<ContractInfoOf<T>>::insert(&self.account_id, info.clone());
@@ -784,10 +784,10 @@ benchmarks! {
 		let code = WasmModule::<T>::from(ModuleDefinition {
 			memory: Some(ImportedMemory::max::<T>()),
 			imported_functions: vec![ImportedFunction {
-				module: "seal0",
+				module: "__unstable__",
 				name: "seal_set_storage",
 				params: vec![ValueType::I32, ValueType::I32, ValueType::I32],
-				return_type: None,
+				return_type: Some(ValueType::I32),
 			}],
 			data_segments: vec![
 				DataSegment {
@@ -800,6 +800,7 @@ benchmarks! {
 				Regular(Instruction::I32Const(0)), // value_ptr
 				Regular(Instruction::I32Const(0)), // value_len
 				Regular(Instruction::Call(0)),
+				Regular(Instruction::Drop),
 			])),
 			.. Default::default()
 		});
@@ -814,10 +815,10 @@ benchmarks! {
 		let code = WasmModule::<T>::from(ModuleDefinition {
 			memory: Some(ImportedMemory::max::<T>()),
 			imported_functions: vec![ImportedFunction {
-				module: "seal0",
+				module: "__unstable__",
 				name: "seal_set_storage",
 				params: vec![ValueType::I32, ValueType::I32, ValueType::I32],
-				return_type: None,
+				return_type: Some(ValueType::I32),
 			}],
 			data_segments: vec![
 				DataSegment {
@@ -830,6 +831,7 @@ benchmarks! {
 				Instruction::I32Const(0), // value_ptr
 				Instruction::I32Const((n * 1024) as i32), // value_len
 				Instruction::Call(0),
+				Instruction::Drop,
 			])),
 			.. Default::default()
 		});
@@ -851,10 +853,10 @@ benchmarks! {
 		let code = WasmModule::<T>::from(ModuleDefinition {
 			memory: Some(ImportedMemory::max::<T>()),
 			imported_functions: vec![ImportedFunction {
-				module: "seal0",
+				module: "__unstable__",
 				name: "seal_clear_storage",
 				params: vec![ValueType::I32],
-				return_type: None,
+				return_type: Some(ValueType::I32),
 			}],
 			data_segments: vec![
 				DataSegment {
@@ -865,6 +867,7 @@ benchmarks! {
 			call_body: Some(body::repeated_dyn(r * API_BENCHMARK_BATCH_SIZE, vec![
 				Counter(0, key_len as u32),
 				Regular(Instruction::Call(0)),
+				Regular(Instruction::Drop),
 			])),
 			.. Default::default()
 		});
@@ -876,6 +879,7 @@ benchmarks! {
 				key.as_slice().try_into().map_err(|e| "Key has wrong length")?,
 				Some(vec![42; T::Schedule::get().limits.payload_len as usize]),
 				None,
+				false,
 			)
 			.map_err(|_| "Failed to write to storage during setup.")?;
 		}
@@ -906,6 +910,10 @@ benchmarks! {
 					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(r * API_BENCHMARK_BATCH_SIZE, vec![
 				Counter(0, key_len as u32), // key_ptr
@@ -924,6 +932,54 @@ benchmarks! {
 				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.")?;
+		}
+		<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 {
+		let r in 0 .. API_BENCHMARK_BATCHES;
+		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_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: "__unstable__",
+				name: "seal_contains_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(r * 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![42; T::Schedule::get().limits.payload_len as usize]),
+				None,
+				false,
 			)
 			.map_err(|_| "Failed to write to storage during setup.")?;
 		}
@@ -970,12 +1026,117 @@ benchmarks! {
 			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![])
 
+	#[skip_meta]
+	seal_take_storage {
+		let r in 0 .. API_BENCHMARK_BATCHES;
+		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_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: "__unstable__",
+				name: "seal_take_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(r * 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![]),
+				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![])
+
+	#[skip_meta]
+	seal_take_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 = sp_std::mem::size_of::<<T::Hashing as sp_runtime::traits::Hash>::Output>();
+		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: "__unstable__",
+				name: "seal_take_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 transfer to unique accounts.
 	seal_transfer {
 		let r in 0 .. API_BENCHMARK_BATCHES;
@@ -2285,7 +2446,7 @@ benchmarks! {
 	// configured `Schedule` during benchmark development.
 	// It can be outputed using the following command:
 	// cargo run --manifest-path=bin/node/cli/Cargo.toml --release \
-	//     --features runtime-benchmarks -- benchmark --dev --execution=native \
+	//     --features runtime-benchmarks -- benchmark --extra --dev --execution=native \
 	//     -p pallet_contracts -e print_schedule --no-median-slopes --no-min-squares
 	#[extra]
 	print_schedule {

substrate/frame/contracts/src/exec.rs

@@ -17,7 +17,7 @@
 
 use crate::{
 	gas::GasMeter,
-	storage::{self, Storage},
+	storage::{self, Storage, WriteOutcome},
 	AccountCounter, BalanceOf, CodeHash, Config, ContractInfo, ContractInfoOf, Error, Event,
 	Pallet as Contracts, Schedule,
 };
@@ -140,9 +140,20 @@ 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 `false` if the `key` wasn't previously set by `set_storage` or
+	/// was deleted.
+	fn contains_storage(&mut self, key: &StorageKey) -> bool;
+
 	/// Sets the storage entry by the given key to the specified value. If `value` is `None` then
 	/// the storage entry is deleted.
-	fn set_storage(&mut self, key: StorageKey, value: Option<Vec<u8>>) -> DispatchResult;
+	fn set_storage(
+		&mut self,
+		key: StorageKey,
+		value: Option<Vec<u8>>,
+		take_old: bool,
+	) -> Result<WriteOutcome, DispatchError>;
 
 	/// Returns a reference to the account id of the caller.
 	fn caller(&self) -> &AccountIdOf<Self::T>;
@@ -985,13 +996,23 @@ where
 		Storage::<T>::read(&self.top_frame_mut().contract_info().trie_id, key)
 	}
 
-	fn set_storage(&mut self, key: StorageKey, value: Option<Vec<u8>>) -> DispatchResult {
+	fn contains_storage(&mut self, key: &StorageKey) -> bool {
+		Storage::<T>::contains(&self.top_frame_mut().contract_info().trie_id, key)
+	}
+
+	fn set_storage(
+		&mut self,
+		key: StorageKey,
+		value: Option<Vec<u8>>,
+		take_old: bool,
+	) -> Result<WriteOutcome, DispatchError> {
 		let frame = self.top_frame_mut();
 		Storage::<T>::write(
 			&frame.contract_info.get(&frame.account_id).trie_id,
 			&key,
 			value,
 			Some(&mut frame.nested_storage),
+			take_old,
 		)
 	}
 
@@ -2349,4 +2370,99 @@ mod tests {
 			assert_eq!(<AccountCounter<Test>>::get(), 4);
 		});
 	}
+
+	#[test]
+	fn set_storage_works() {
+		let code_hash = MockLoader::insert(Call, |ctx, _| {
+			// Write
+			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![4, 5, 6]), true),
+				Ok(WriteOutcome::New)
+			);
+			assert_eq!(ctx.ext.set_storage([3; 32], None, false), Ok(WriteOutcome::New));
+			assert_eq!(ctx.ext.set_storage([4; 32], None, true), Ok(WriteOutcome::New));
+			assert_eq!(ctx.ext.set_storage([5; 32], Some(vec![]), false), Ok(WriteOutcome::New));
+			assert_eq!(ctx.ext.set_storage([6; 32], Some(vec![]), true), Ok(WriteOutcome::New));
+
+			// Overwrite
+			assert_eq!(
+				ctx.ext.set_storage([1; 32], Some(vec![42]), false),
+				Ok(WriteOutcome::Overwritten(3))
+			);
+			assert_eq!(
+				ctx.ext.set_storage([2; 32], Some(vec![48]), true),
+				Ok(WriteOutcome::Taken(vec![4, 5, 6]))
+			);
+			assert_eq!(ctx.ext.set_storage([3; 32], None, false), Ok(WriteOutcome::New));
+			assert_eq!(ctx.ext.set_storage([4; 32], None, true), Ok(WriteOutcome::New));
+			assert_eq!(
+				ctx.ext.set_storage([5; 32], Some(vec![]), false),
+				Ok(WriteOutcome::Overwritten(0))
+			);
+			assert_eq!(
+				ctx.ext.set_storage([6; 32], Some(vec![]), true),
+				Ok(WriteOutcome::Taken(vec![]))
+			);
+
+			exec_success()
+		});
+
+		ExtBuilder::default().build().execute_with(|| {
+			let min_balance = <Test as Config>::Currency::minimum_balance();
+			let schedule = <Test as Config>::Schedule::get();
+			let mut gas_meter = GasMeter::<Test>::new(GAS_LIMIT);
+			set_balance(&ALICE, min_balance * 1000);
+			place_contract(&BOB, code_hash);
+			let mut storage_meter = storage::meter::Meter::new(&ALICE, None, 0).unwrap();
+			assert_ok!(MockStack::run_call(
+				ALICE,
+				BOB,
+				&mut gas_meter,
+				&mut storage_meter,
+				&schedule,
+				0,
+				vec![],
+				None,
+			));
+		});
+	}
+
+	#[test]
+	fn contains_storage_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.contains_storage(&[1; 32]), true);
+			assert_eq!(ctx.ext.contains_storage(&[3; 32]), false);
+
+			exec_success()
+		});
+
+		ExtBuilder::default().build().execute_with(|| {
+			let min_balance = <Test as Config>::Currency::minimum_balance();
+			let schedule = <Test as Config>::Schedule::get();
+			let mut gas_meter = GasMeter::<Test>::new(GAS_LIMIT);
+			set_balance(&ALICE, min_balance * 1000);
+			place_contract(&BOB, code_hash);
+			let mut storage_meter = storage::meter::Meter::new(&ALICE, None, 0).unwrap();
+			assert_ok!(MockStack::run_call(
+				ALICE,
+				BOB,
+				&mut gas_meter,
+				&mut storage_meter,
+				&schedule,
+				0,
+				vec![],
+				None,
+			));
+		});
+	}
 }

substrate/frame/contracts/src/schedule.rs

@@ -322,12 +322,21 @@ pub struct HostFnWeights<T: Config> {
 	/// Weight of calling `seal_clear_storage`.
 	pub clear_storage: Weight,
 
+	/// Weight of calling `seal_contains_storage`.
+	pub contains_storage: Weight,
+
 	/// Weight of calling `seal_get_storage`.
 	pub get_storage: Weight,
 
 	/// Weight per byte of an item received via `seal_get_storage`.
 	pub get_storage_per_byte: Weight,
 
+	/// Weight of calling `seal_take_storage`.
+	pub take_storage: Weight,
+
+	/// Weight per byte of an item received via `seal_take_storage`.
+	pub take_storage_per_byte: Weight,
+
 	/// Weight of calling `seal_transfer`.
 	pub transfer: Weight,
 
@@ -579,8 +588,11 @@ impl<T: Config> Default for HostFnWeights<T> {
 			set_storage: cost_batched!(seal_set_storage),
 			set_storage_per_byte: cost_byte_batched!(seal_set_storage_per_kb),
 			clear_storage: cost_batched!(seal_clear_storage),
+			contains_storage: cost_batched!(seal_contains_storage),
 			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),
+			take_storage_per_byte: cost_byte_batched!(seal_take_storage_per_kb),
 			transfer: cost_batched!(seal_transfer),
 			call: cost_batched!(seal_call),
 			call_transfer_surcharge: cost_batched_args!(

substrate/frame/contracts/src/storage.rs

@@ -72,6 +72,21 @@ pub struct DeletedContract {
 	pub(crate) trie_id: TrieId,
 }
 
+/// Information about what happended to the pre-existing value when calling [`Storage::write`].
+#[cfg_attr(test, derive(Debug, PartialEq))]
+pub enum WriteOutcome {
+	/// No value existed at the specified key.
+	New,
+	/// A value of the returned length was overwritten.
+	Overwritten(u32),
+	/// The returned value was taken out of storage before being overwritten.
+	///
+	/// This is only returned when specifically requested because it causes additional work
+	/// depending on the size of the pre-existing value. When not requested [`Self::Overwritten`]
+	/// is returned instead.
+	Taken(Vec<u8>),
+}
+
 pub struct Storage<T>(PhantomData<T>);
 
 impl<T> Storage<T>
@@ -87,9 +102,15 @@ where
 		child::get_raw(&child_trie_info(trie_id), &blake2_256(key))
 	}
 
+	/// Returns `true` iff the `key` exists in storage.
+	pub fn contains(trie_id: &TrieId, key: &StorageKey) -> bool {
+		child::exists(&child_trie_info(trie_id), &blake2_256(key))
+	}
+
 	/// Update a storage entry into a contract's kv storage.
 	///
-	/// If the `new_value` is `None` then the kv pair is removed.
+	/// If the `new_value` is `None` then the kv pair is removed. If `take` is true
+	/// a [`WriteOutcome::Taken`] is returned instead of a [`WriteOutcome::Overwritten`].
 	///
 	/// This function also records how much storage was created or removed if a `storage_meter`
 	/// is supplied. It should only be absent for testing or benchmarking code.
@@ -98,13 +119,19 @@ where
 		key: &StorageKey,
 		new_value: Option<Vec<u8>>,
 		storage_meter: Option<&mut meter::NestedMeter<T>>,
-	) -> DispatchResult {
+		take: bool,
+	) -> Result<WriteOutcome, DispatchError> {
 		let hashed_key = blake2_256(key);
 		let child_trie_info = &child_trie_info(trie_id);
+		let (old_len, old_value) = if take {
+			let val = child::get_raw(&child_trie_info, &hashed_key);
+			(val.as_ref().map(|v| v.len() as u32), val)
+		} else {
+			(child::len(&child_trie_info, &hashed_key), None)
+		};
 
 		if let Some(storage_meter) = storage_meter {
 			let mut diff = meter::Diff::default();
-			let old_len = child::len(&child_trie_info, &hashed_key);
 			match (old_len, new_value.as_ref().map(|v| v.len() as u32)) {
 				(Some(old_len), Some(new_len)) =>
 					if new_len > old_len {
@@ -130,7 +157,11 @@ where
 			None => child::kill(&child_trie_info, &hashed_key),
 		}
 
-		Ok(())
+		Ok(match (old_len, old_value) {
+			(None, _) => WriteOutcome::New,
+			(Some(old_len), None) => WriteOutcome::Overwritten(old_len),
+			(Some(_), Some(old_value)) => WriteOutcome::Taken(old_value),
+		})
 	}
 
 	/// Creates a new contract descriptor in the storage with the given code hash at the given

substrate/frame/contracts/src/tests.rs

@@ -1545,7 +1545,8 @@ fn lazy_removal_partial_remove_works() {
 
 		// Put value into the contracts child trie
 		for val in &vals {
-			Storage::<Test>::write(&info.trie_id, &val.0, Some(val.2.clone()), None).unwrap();
+			Storage::<Test>::write(&info.trie_id, &val.0, Some(val.2.clone()), None, false)
+				.unwrap();
 		}
 		<ContractInfoOf<Test>>::insert(&addr, info.clone());
 
@@ -1629,7 +1630,8 @@ fn lazy_removal_does_no_run_on_full_block() {
 
 		// Put value into the contracts child trie
 		for val in &vals {
-			Storage::<Test>::write(&info.trie_id, &val.0, Some(val.2.clone()), None).unwrap();
+			Storage::<Test>::write(&info.trie_id, &val.0, Some(val.2.clone()), None, false)
+				.unwrap();
 		}
 		<ContractInfoOf<Test>>::insert(&addr, info.clone());
 
@@ -1712,7 +1714,8 @@ fn lazy_removal_does_not_use_all_weight() {
 
 		// Put value into the contracts child trie
 		for val in &vals {
-			Storage::<Test>::write(&info.trie_id, &val.0, Some(val.2.clone()), None).unwrap();
+			Storage::<Test>::write(&info.trie_id, &val.0, Some(val.2.clone()), None, false)
+				.unwrap();
 		}
 		<ContractInfoOf<Test>>::insert(&addr, info.clone());
 

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

@@ -257,21 +257,22 @@ mod tests {
 			AccountIdOf, BlockNumberOf, ErrorOrigin, ExecError, Executable, Ext, SeedOf, StorageKey,
 		},
 		gas::GasMeter,
+		storage::WriteOutcome,
 		tests::{Call, Test, ALICE, BOB},
 		BalanceOf, CodeHash, Error, Pallet as Contracts,
 	};
 	use assert_matches::assert_matches;
-	use frame_support::{
-		assert_ok,
-		dispatch::{DispatchResult, DispatchResultWithPostInfo},
-		weights::Weight,
-	};
+	use frame_support::{assert_ok, dispatch::DispatchResultWithPostInfo, weights::Weight};
 	use hex_literal::hex;
 	use pallet_contracts_primitives::{ExecReturnValue, ReturnFlags};
 	use pretty_assertions::assert_eq;
 	use sp_core::{Bytes, H256};
 	use sp_runtime::DispatchError;
-	use std::{borrow::BorrowMut, cell::RefCell, collections::HashMap};
+	use std::{
+		borrow::BorrowMut,
+		cell::RefCell,
+		collections::hash_map::{Entry, HashMap},
+	};
 
 	#[derive(Debug, PartialEq, Eq)]
 	struct InstantiateEntry {
@@ -384,9 +385,26 @@ mod tests {
 		fn get_storage(&mut self, key: &StorageKey) -> Option<Vec<u8>> {
 			self.storage.get(key).cloned()
 		}
-		fn set_storage(&mut self, key: StorageKey, value: Option<Vec<u8>>) -> DispatchResult {
-			*self.storage.entry(key).or_insert(Vec::new()) = value.unwrap_or(Vec::new());
-			Ok(())
+		fn contains_storage(&mut self, key: &StorageKey) -> bool {
+			self.storage.contains_key(key)
+		}
+		fn set_storage(
+			&mut self,
+			key: StorageKey,
+			value: Option<Vec<u8>>,
+			take_old: bool,
+		) -> Result<WriteOutcome, DispatchError> {
+			let entry = self.storage.entry(key);
+			let result = match (entry, take_old) {
+				(Entry::Vacant(_), _) => WriteOutcome::New,
+				(Entry::Occupied(entry), false) =>
+					WriteOutcome::Overwritten(entry.remove().len() as u32),
+				(Entry::Occupied(entry), true) => WriteOutcome::Taken(entry.remove()),
+			};
+			if let Some(value) = value {
+				self.storage.insert(key, value);
+			}
+			Ok(result)
 		}
 		fn caller(&self) -> &AccountIdOf<Self::T> {
 			&ALICE
@@ -932,8 +950,8 @@ mod tests {
 		"\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11"
 	)
 
-	;; [32, 36) buffer size = 128 bytes
-	(data (i32.const 32) "\80")
+	;; [32, 36) buffer size = 4k in little endian
+	(data (i32.const 32) "\00\10")
 
 	;; [36; inf) buffer where the result is copied
 
@@ -1957,4 +1975,278 @@ mod tests {
 		);
 		assert_eq!(*ext.runtime_calls.borrow(), vec![]);
 	}
+
+	#[test]
+	#[cfg(feature = "unstable-interface")]
+	fn set_storage_works() {
+		const CODE: &str = r#"
+(module
+	(import "seal0" "seal_input" (func $seal_input (param i32 i32)))
+	(import "seal0" "seal_return" (func $seal_return (param i32 i32 i32)))
+	(import "__unstable__" "seal_set_storage" (func $seal_set_storage (param i32 i32 i32) (result i32)))
+	(import "env" "memory" (memory 1 1))
+
+	;; 0x1000 = 4k in little endian
+	;; size of input buffer
+	(data (i32.const 0) "\00\10")
+
+	(func (export "call")
+		;; Receive (key ++ value_to_write)
+		(call $seal_input
+			(i32.const 4)	;; Pointer to the input buffer
+			(i32.const 0)	;; Size of the length buffer
+		)
+		;; Store the passed value to the passed key and store result to memory
+		(i32.store (i32.const 0)
+			(call $seal_set_storage
+				(i32.const 4)				;; key_ptr
+				(i32.const 36)				;; value_ptr
+				(i32.sub					;; value_len (input_size - key_size)
+					(i32.load (i32.const 0))
+					(i32.const 32)
+				)
+			)
+		)
+		(call $seal_return
+			(i32.const 0)	;; flags
+			(i32.const 0)	;; returned value
+			(i32.const 4)	;; length of returned value
+		)
+	)
+
+	(func (export "deploy"))
+)
+"#;
+
+		let mut ext = MockExt::default();
+
+		// 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!(ext.storage.get(&[1u8; 32]).unwrap(), &[42u8, 48]);
+
+		// value do exist -> length of old value returned
+		let input = ([1u8; 32], [0u8; 0]).encode();
+		let result = execute(CODE, input, &mut ext).unwrap();
+		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), 2);
+		assert_eq!(ext.storage.get(&[1u8; 32]).unwrap(), &[0u8; 0]);
+
+		// value do exist -> length of old value returned (test for zero sized val)
+		let input = ([1u8; 32], [99u8]).encode();
+		let result = execute(CODE, input, &mut ext).unwrap();
+		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), 0);
+		assert_eq!(ext.storage.get(&[1u8; 32]).unwrap(), &[99u8]);
+	}
+
+	#[test]
+	#[cfg(feature = "unstable-interface")]
+	fn clear_storage_works() {
+		const CODE: &str = r#"
+(module
+	(import "seal0" "seal_input" (func $seal_input (param i32 i32)))
+	(import "seal0" "seal_return" (func $seal_return (param i32 i32 i32)))
+	(import "__unstable__" "seal_clear_storage" (func $seal_clear_storage (param i32) (result i32)))
+	(import "env" "memory" (memory 1 1))
+
+	;; 0x1000 = 4k in little endian
+	;; size of input buffer
+	(data (i32.const 0) "\00\10")
+
+	(func (export "call")
+		;; Receive key
+		(call $seal_input
+			(i32.const 4)	;; Pointer to the input buffer
+			(i32.const 0)	;; Size of the length buffer
+		)
+		;; Store the passed value to the passed key and store result to memory
+		(i32.store (i32.const 0)
+			(call $seal_clear_storage
+				(i32.const 4)				;; key_ptr
+			)
+		)
+		(call $seal_return
+			(i32.const 0)	;; flags
+			(i32.const 0)	;; returned value
+			(i32.const 4)	;; length of returned value
+		)
+	)
+
+	(func (export "deploy"))
+)
+"#;
+
+		let mut ext = MockExt::default();
+
+		ext.storage.insert([1u8; 32], vec![42u8]);
+		ext.storage.insert([2u8; 32], vec![]);
+
+		// 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!(ext.storage.get(&[3u8; 32]), None);
+
+		// value did exist -> length returned
+		let result = execute(CODE, [1u8; 32].encode(), &mut ext).unwrap();
+		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), 1);
+		assert_eq!(ext.storage.get(&[1u8; 32]), None);
+
+		// value did exist -> length returned (test for 0 sized)
+		let result = execute(CODE, [2u8; 32].encode(), &mut ext).unwrap();
+		assert_eq!(u32::from_le_bytes(result.data.0.try_into().unwrap()), 0);
+		assert_eq!(ext.storage.get(&[2u8; 32]), None);
+	}
+
+	#[test]
+	#[cfg(feature = "unstable-interface")]
+	fn take_storage_works() {
+		const CODE: &str = r#"
+(module
+	(import "seal0" "seal_return" (func $seal_return (param i32 i32 i32)))
+	(import "seal0" "seal_input" (func $seal_input (param i32 i32)))
+	(import "__unstable__" "seal_take_storage" (func $seal_take_storage (param i32 i32 i32) (result i32)))
+	(import "env" "memory" (memory 1 1))
+
+	;; [0, 32) size of input buffer (32 byte as we copy the key here)
+	(data (i32.const 0) "\20")
+
+	;; [32, 64) size of output buffer
+	;; 4k in little endian
+	(data (i32.const 32) "\00\10")
+
+	;; [64, 96) input buffer
+
+	;; [96, inf) output buffer
+
+	(func (export "call")
+		;; Receive key
+		(call $seal_input
+			(i32.const 64)	;; Pointer to the input buffer
+			(i32.const 0)	;; Size of the length buffer
+		)
+
+		;; Load a storage value and result of this call into the output buffer
+		(i32.store (i32.const 96)
+			(call $seal_take_storage
+				(i32.const 64)		;; The pointer to the storage key to fetch
+				(i32.const 100)		;; Pointer to the output buffer
+				(i32.const 32)		;; Pointer to the size of the buffer
+			)
+		)
+
+		;; Return the contents of the buffer
+		(call $seal_return
+			(i32.const 0)					;; flags
+			(i32.const 96)					;; output buffer ptr
+			(i32.add						;; length: storage size + 4 (retval)
+				(i32.load (i32.const 32))
+				(i32.const 4)
+			)
+		)
+	)
+
+	(func (export "deploy"))
+)
+"#;
+
+		let mut ext = MockExt::default();
+
+		ext.storage.insert([1u8; 32], vec![42u8]);
+		ext.storage.insert([2u8; 32], vec![]);
+
+		// value does not exist -> error returned
+		let result = execute(CODE, [3u8; 32].encode(), &mut ext).unwrap();
+		assert_eq!(
+			u32::from_le_bytes(result.data.0[0..4].try_into().unwrap()),
+			ReturnCode::KeyNotFound as u32
+		);
+
+		// value did exist -> value returned
+		let result = execute(CODE, [1u8; 32].encode(), &mut ext).unwrap();
+		assert_eq!(
+			u32::from_le_bytes(result.data.0[0..4].try_into().unwrap()),
+			ReturnCode::Success as u32
+		);
+		assert_eq!(ext.storage.get(&[1u8; 32]), None);
+		assert_eq!(&result.data.0[4..], &[42u8]);
+
+		// value did exist -> length returned (test for 0 sized)
+		let result = execute(CODE, [2u8; 32].encode(), &mut ext).unwrap();
+		assert_eq!(
+			u32::from_le_bytes(result.data.0[0..4].try_into().unwrap()),
+			ReturnCode::Success as u32
+		);
+		assert_eq!(ext.storage.get(&[2u8; 32]), None);
+		assert_eq!(&result.data.0[4..], &[0u8; 0]);
+	}
+
+	#[test]
+	#[cfg(feature = "unstable-interface")]
+	fn contains_storage_works() {
+		const CODE: &str = r#"
+(module
+	(import "seal0" "seal_return" (func $seal_return (param i32 i32 i32)))
+	(import "seal0" "seal_input" (func $seal_input (param i32 i32)))
+	(import "__unstable__" "seal_contains_storage" (func $seal_contains_storage (param i32) (result i32)))
+	(import "env" "memory" (memory 1 1))
+
+	;; [0, 4) size of input buffer (32 byte as we copy the key here)
+	(data (i32.const 0) "\20")
+
+	;; [4, 36) input buffer
+
+	;; [36, inf) output buffer
+
+	(func (export "call")
+		;; Receive key
+		(call $seal_input
+			(i32.const 4)	;; Pointer to the input buffer
+			(i32.const 0)	;; Size of the length buffer
+		)
+
+		;; Load the return value into the output buffer
+		(i32.store (i32.const 36)
+			(call $seal_contains_storage
+				(i32.const 4)		;; The pointer to the storage key to fetch
+			)
+		)
+
+		;; Return the contents of the buffer
+		(call $seal_return
+			(i32.const 0)					;; flags
+			(i32.const 36)					;; output buffer ptr
+			(i32.const 4)					;; result is integer (4 bytes)
+		)
+	)
+
+	(func (export "deploy"))
+)
+"#;
+
+		let mut ext = MockExt::default();
+
+		ext.storage.insert([1u8; 32], vec![42u8]);
+		ext.storage.insert([2u8; 32], vec![]);
+
+		// value does not exist -> error returned
+		let result = execute(CODE, [3u8; 32].encode(), &mut ext).unwrap();
+		assert_eq!(
+			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!(
+			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!(
+			u32::from_le_bytes(result.data.0.try_into().unwrap()),
+			ReturnCode::Success as u32
+		);
+	}
 }

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

@@ -21,6 +21,7 @@ use crate::{
 	exec::{ExecError, ExecResult, Ext, StorageKey, TopicOf},
 	gas::{ChargedAmount, Token},
 	schedule::HostFnWeights,
+	storage::WriteOutcome,
 	wasm::env_def::ConvertibleToWasm,
 	BalanceOf, CodeHash, Config, Error,
 };
@@ -172,10 +173,19 @@ pub enum RuntimeCosts {
 	SetStorage(u32),
 	/// Weight of calling `seal_clear_storage`.
 	ClearStorage,
+	/// Weight of calling `seal_contains_storage`.
+	#[cfg(feature = "unstable-interface")]
+	ContainsStorage,
 	/// Weight of calling `seal_get_storage` without output weight.
 	GetStorageBase,
 	/// Weight of an item received via `seal_get_storage` for the given size.
 	GetStorageCopyOut(u32),
+	/// Weight of calling `seal_take_storage` without output weight.
+	#[cfg(feature = "unstable-interface")]
+	TakeStorageBase,
+	/// 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.
@@ -242,8 +252,14 @@ impl RuntimeCosts {
 			SetStorage(len) =>
 				s.set_storage.saturating_add(s.set_storage_per_byte.saturating_mul(len.into())),
 			ClearStorage => s.clear_storage,
+			#[cfg(feature = "unstable-interface")]
+			ContainsStorage => s.contains_storage,
 			GetStorageBase => s.get_storage,
 			GetStorageCopyOut(len) => s.get_storage_per_byte.saturating_mul(len.into()),
+			#[cfg(feature = "unstable-interface")]
+			TakeStorageBase => s.take_storage,
+			#[cfg(feature = "unstable-interface")]
+			TakeStorageCopyOut(len) => 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())),
@@ -632,6 +648,50 @@ 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))?;
+		if value_len > self.ext.max_value_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
+			.set_storage(key, value, false)
+			.map(Self::overwritten_len)
+			.map_err(Into::into)
+	}
+
+	fn clear_storage(&mut self, key_ptr: u32) -> Result<u32, TrapReason> {
+		self.charge_gas(RuntimeCosts::ClearStorage)?;
+		let mut key: StorageKey = [0; 32];
+		self.read_sandbox_memory_into_buf(key_ptr, &mut key)?;
+		self.ext
+			.set_storage(key, None, false)
+			.map(Self::overwritten_len)
+			.map_err(Into::into)
+	}
+
 	fn call(
 		&mut self,
 		flags: CallFlags,
@@ -745,10 +805,18 @@ define_env!(Env, <E: Ext>,
 		Ok(())
 	},
 
+	// Set the value at the given key in the contract storage.
+	//
+	// Equivalent to the newer version of `seal_set_storage` with the exception of the return
+	// type. Still a valid thing to call when not interested in the return value.
+	[seal0] seal_set_storage(ctx, key_ptr: u32, value_ptr: u32, value_len: u32) => {
+		ctx.set_storage(key_ptr, value_ptr, value_len).map(|_| ())
+	},
+
 	// Set the value at the given key in the contract storage.
 	//
 	// The value length must not exceed the maximum defined by the contracts module parameters.
-	// Storing an empty value is disallowed.
+	// Specifying a `value_len` of zero will store an empty value.
 	//
 	// # Parameters
 	//
@@ -756,19 +824,20 @@ define_env!(Env, <E: Ext>,
 	// - `value_ptr`: pointer into the linear memory where the value to set is placed.
 	// - `value_len`: the length of the value in bytes.
 	//
-	// # Traps
+	// # Return Value
 	//
-	// - If value length exceeds the configured maximum value length of a storage entry.
-	// - Upon trying to set an empty storage entry (value length is 0).
-	[seal0] seal_set_storage(ctx, key_ptr: u32, value_ptr: u32, value_len: u32) => {
-		ctx.charge_gas(RuntimeCosts::SetStorage(value_len))?;
-		if value_len > ctx.ext.max_value_size() {
-			Err(Error::<E::T>::ValueTooLarge)?;
-		}
-		let mut key: StorageKey = [0; 32];
-		ctx.read_sandbox_memory_into_buf(key_ptr, &mut key)?;
-		let value = Some(ctx.read_sandbox_memory(value_ptr, value_len)?);
-		ctx.ext.set_storage(key, value).map_err(Into::into)
+	// Returns the size of the pre-existing value at the specified key if any. Otherwise
+	// `u32::MAX` 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)
+	},
+
+	// Clear the value at the given key in the contract storage.
+	//
+	// Equivalent to the newer version of `seal_clear_storage` with the exception of the return
+	// type. Still a valid thing to call when not interested in the return value.
+	[seal0] seal_clear_storage(ctx, key_ptr: u32) => {
+		ctx.clear_storage(key_ptr).map(|_| ()).map_err(Into::into)
 	},
 
 	// Clear the value at the given key in the contract storage.
@@ -776,11 +845,13 @@ define_env!(Env, <E: Ext>,
 	// # Parameters
 	//
 	// - `key_ptr`: pointer into the linear memory where the location to clear the value is placed.
-	[seal0] seal_clear_storage(ctx, key_ptr: u32) => {
-		ctx.charge_gas(RuntimeCosts::ClearStorage)?;
-		let mut key: StorageKey = [0; 32];
-		ctx.read_sandbox_memory_into_buf(key_ptr, &mut key)?;
-		ctx.ext.set_storage(key, None).map_err(Into::into)
+	//
+	// # 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.
+	[__unstable__] seal_clear_storage(ctx, key_ptr: u32) -> u32 => {
+		ctx.clear_storage(key_ptr).map_err(Into::into)
 	},
 
 	// Retrieve the value under the given key from storage.
@@ -809,6 +880,55 @@ define_env!(Env, <E: Ext>,
 		}
 	},
 
+	// 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
+	//
+	// `ReturnCode::KeyNotFound`
+	[__unstable__] seal_contains_storage(ctx, key_ptr: u32) -> ReturnCode => {
+		ctx.charge_gas(RuntimeCosts::ContainsStorage)?;
+		let mut key: StorageKey = [0; 32];
+		ctx.read_sandbox_memory_into_buf(key_ptr, &mut key)?;
+		if ctx.ext.contains_storage(&key) {
+			Ok(ReturnCode::Success)
+		} else {
+			Ok(ReturnCode::KeyNotFound)
+		}
+	},
+
+	// Retrieve and remove the value under the given key from storage.
+	//
+	// # Parameters
+	//
+	// - `key_ptr`: pointer into the linear memory where the key of the requested value is placed.
+	// - `out_ptr`: pointer to the linear memory where the value is written to.
+	// - `out_len_ptr`: in-out pointer into linear memory where the buffer length
+	//   is read from and the value length is written to.
+	//
+	// # Errors
+	//
+	// `ReturnCode::KeyNotFound`
+	[__unstable__] seal_take_storage(ctx, key_ptr: u32, out_ptr: u32, out_len_ptr: u32) -> ReturnCode => {
+		ctx.charge_gas(RuntimeCosts::TakeStorageBase)?;
+		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)? {
+			ctx.write_sandbox_output(out_ptr, out_len_ptr, &value, false, |len| {
+				Some(RuntimeCosts::TakeStorageCopyOut(len))
+			})?;
+			Ok(ReturnCode::Success)
+		} else {
+			Ok(ReturnCode::KeyNotFound)
+		}
+	},
+
 	// Transfer some value to another account.
 	//
 	// # Parameters