Add new versions for storage access host functions

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,7 @@ 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.")?;
 		}
@@ -970,12 +979,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;

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,
 };
@@ -142,7 +142,12 @@ pub trait Ext: sealing::Sealed {
 
 	/// 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 +990,19 @@ 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 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 +2360,64 @@ 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,
+			));
+		});
+	}
 }

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>
@@ -89,7 +104,8 @@ where
 
 	/// 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 +114,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 +152,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