revive: prototype new storage host fn to handle hot vs cold load

polkadot/node/core/pvf/common/src/executor_interface.rs

@@ -219,6 +219,10 @@ impl sp_externalities::Externalities for ValidationExternalities {
 		panic!("storage: unsupported feature for parachain validation")
 	}
 
+	fn storage_with_status(&mut self, _: &[u8]) -> sp_externalities::StateLoad<Option<Vec<u8>>> {
+		panic!("storage_with_status: unsupported feature for parachain validation")
+	}
+
 	fn storage_hash(&mut self, _: &[u8]) -> Option<Vec<u8>> {
 		panic!("storage_hash: unsupported feature for parachain validation")
 	}
@@ -231,6 +235,14 @@ impl sp_externalities::Externalities for ValidationExternalities {
 		panic!("child_storage: unsupported feature for parachain validation")
 	}
 
+	fn child_storage_with_status(
+		&mut self,
+		_: &ChildInfo,
+		_: &[u8],
+	) -> sp_externalities::StateLoad<Option<Vec<u8>>> {
+		panic!("child_storage_with_status: unsupported feature for parachain validation")
+	}
+
 	fn kill_child_storage(
 		&mut self,
 		_child_info: &ChildInfo,

substrate/client/db/Cargo.toml

@@ -40,6 +40,7 @@ sp-arithmetic = { workspace = true, default-features = true }
 sp-blockchain = { workspace = true, default-features = true }
 sp-core = { workspace = true, default-features = true }
 sp-database = { workspace = true, default-features = true }
+sp-externalities = { workspace = true, default-features = true }
 sp-runtime = { workspace = true, default-features = true }
 sp-state-machine = { workspace = true, default-features = true }
 sp-trie = { workspace = true, default-features = true }

substrate/client/db/src/bench.rs

@@ -353,6 +353,15 @@ impl<Hasher: Hash> StateBackend<Hasher> for BenchmarkingState<Hasher> {
 		self.state.borrow().as_ref().ok_or_else(state_err)?.storage(key)
 	}
 
+	fn storage_with_status(
+		&self,
+		key: &[u8],
+	) -> Result<sp_externalities::StateLoad<Option<Vec<u8>>>, Self::Error> {
+		// Only if cold load ?
+		self.add_read_key(None, key);
+		self.state.borrow().as_ref().ok_or_else(state_err)?.storage_with_status(key)
+	}
+
 	fn storage_hash(&self, key: &[u8]) -> Result<Option<Hasher::Output>, Self::Error> {
 		self.add_read_key(None, key);
 		self.state.borrow().as_ref().ok_or_else(state_err)?.storage_hash(key)
@@ -371,6 +380,20 @@ impl<Hasher: Hash> StateBackend<Hasher> for BenchmarkingState<Hasher> {
 			.child_storage(child_info, key)
 	}
 
+	fn child_storage_with_status(
+		&self,
+		child_info: &ChildInfo,
+		key: &[u8],
+	) -> Result<sp_externalities::StateLoad<Option<Vec<u8>>>, Self::Error> {
+		// Only if cold load ?
+		self.add_read_key(Some(child_info.storage_key()), key);
+		self.state
+			.borrow()
+			.as_ref()
+			.ok_or_else(state_err)?
+			.child_storage_with_status(child_info, key)
+	}
+
 	fn child_storage_hash(
 		&self,
 		child_info: &ChildInfo,

substrate/client/db/src/lib.rs

@@ -78,6 +78,7 @@ use sp_core::{
 	storage::{well_known_keys, ChildInfo},
 };
 use sp_database::Transaction;
+use sp_externalities::StateLoad;
 use sp_runtime::{
 	generic::BlockId,
 	traits::{
@@ -199,6 +200,10 @@ impl<B: BlockT> StateBackend<HashingFor<B>> for RefTrackingState<B> {
 		self.state.storage(key)
 	}
 
+	fn storage_with_status(&self, key: &[u8]) -> Result<StateLoad<Option<Vec<u8>>>, Self::Error> {
+		self.state.storage_with_status(key)
+	}
+
 	fn storage_hash(&self, key: &[u8]) -> Result<Option<B::Hash>, Self::Error> {
 		self.state.storage_hash(key)
 	}
@@ -211,6 +216,14 @@ impl<B: BlockT> StateBackend<HashingFor<B>> for RefTrackingState<B> {
 		self.state.child_storage(child_info, key)
 	}
 
+	fn child_storage_with_status(
+		&self,
+		child_info: &ChildInfo,
+		key: &[u8],
+	) -> Result<StateLoad<Option<Vec<u8>>>, Self::Error> {
+		self.state.child_storage_with_status(child_info, key)
+	}
+
 	fn child_storage_hash(
 		&self,
 		child_info: &ChildInfo,

substrate/client/db/src/record_stats_state.rs

@@ -20,6 +20,7 @@
 
 use crate::stats::StateUsageStats;
 use sp_core::storage::ChildInfo;
+use sp_externalities::StateLoad;
 use sp_runtime::{
 	traits::{Block as BlockT, HashingFor},
 	StateVersion,
@@ -119,6 +120,12 @@ impl<S: StateBackend<HashingFor<B>>, B: BlockT> StateBackend<HashingFor<B>>
 		Ok(value)
 	}
 
+	fn storage_with_status(&self, key: &[u8]) -> Result<StateLoad<Option<Vec<u8>>>, Self::Error> {
+		let result = self.state.storage_with_status(key)?;
+		self.usage.tally_key_read(key, result.data.as_ref(), false);
+		Ok(result)
+	}
+
 	fn storage_hash(&self, key: &[u8]) -> Result<Option<B::Hash>, Self::Error> {
 		self.state.storage_hash(key)
 	}
@@ -137,6 +144,20 @@ impl<S: StateBackend<HashingFor<B>>, B: BlockT> StateBackend<HashingFor<B>>
 		Ok(value)
 	}
 
+	fn child_storage_with_status(
+		&self,
+		child_info: &ChildInfo,
+		key: &[u8],
+	) -> Result<StateLoad<Option<Vec<u8>>>, Self::Error> {
+		let key_tuple = (child_info.storage_key().to_vec(), key.to_vec());
+		let result = self.state.child_storage_with_status(child_info, &key_tuple.1)?;
+
+		// just pass it through the usage counter
+		let value = self.usage.tally_child_key_read(&key_tuple, result.data, false);
+
+		Ok(StateLoad { is_cold: result.is_cold, data: value })
+	}
+
 	fn child_storage_hash(
 		&self,
 		child_info: &ChildInfo,

substrate/frame/revive/src/benchmarking.rs

@@ -1407,10 +1407,12 @@ mod benchmarks {
 
 	// n: new byte size
 	// o: old byte size
+	// c: is_cold (0 = hot, 1 = cold)
 	#[benchmark(skip_meta, pov_mode = Measured)]
 	fn seal_set_storage(
 		n: Linear<0, { limits::STORAGE_BYTES }>,
 		o: Linear<0, { limits::STORAGE_BYTES }>,
+		c: Linear<0, 1>,
 	) -> Result<(), BenchmarkError> {
 		let max_key_len = limits::STORAGE_KEY_BYTES;
 		let key = Key::try_from_var(vec![0u8; max_key_len as usize])
@@ -1423,6 +1425,13 @@ mod benchmarks {
 		info.write(&key, Some(vec![42u8; o as usize]), None, false)
 			.map_err(|_| "Failed to write to storage during setup.")?;
 
+		// Whitelist key if c=0 (hot)
+		if c == 0 {
+			let mut full_key = info.child_trie_info().prefixed_storage_key().into_inner();
+			full_key.extend_from_slice(&key.hash());
+			frame_benchmarking::benchmarking::add_to_whitelist(full_key.into());
+		}
+
 		let result;
 		#[block]
 		{
@@ -1437,12 +1446,15 @@ mod benchmarks {
 		}
 
 		assert_ok!(result);
-		assert_eq!(info.read(&key).unwrap(), value);
+		assert_eq!(info.read(&key).data.unwrap(), value);
 		Ok(())
 	}
 
 	#[benchmark(skip_meta, pov_mode = Measured)]
-	fn clear_storage(n: Linear<0, { limits::STORAGE_BYTES }>) -> Result<(), BenchmarkError> {
+	fn clear_storage(
+		n: Linear<0, { limits::STORAGE_BYTES }>,
+		c: Linear<0, 1>,
+	) -> Result<(), BenchmarkError> {
 		let max_key_len = limits::STORAGE_KEY_BYTES;
 		let key = Key::try_from_var(vec![0u8; max_key_len as usize])
 			.map_err(|_| "Key has wrong length")?;
@@ -1455,6 +1467,15 @@ mod benchmarks {
 		.abi_encode();
 
 		let mut call_setup = CallSetup::<T>::default();
+
+		// Whitelist key if c=0 (hot)
+		if c == 0 {
+			let info = call_setup.contract().info()?;
+			let mut full_key = info.child_trie_info().prefixed_storage_key().into_inner();
+			full_key.extend_from_slice(&key.hash());
+			frame_benchmarking::benchmarking::add_to_whitelist(full_key.into());
+		}
+
 		let (mut ext, _) = call_setup.ext();
 		ext.set_storage(&key, Some(vec![42u8; max_key_len as usize]), false)
 			.map_err(|_| "Failed to write to storage during setup.")?;
@@ -1465,17 +1486,21 @@ mod benchmarks {
 			result = run_builtin_precompile(
 				&mut ext,
 				H160(BenchmarkStorage::<T>::MATCHER.base_address()).as_fixed_bytes(),
-				input_bytes,
+				input_bytes.clone(),
 			);
 		}
 		assert_ok!(result);
-		assert!(ext.get_storage(&key).is_none());
+		assert!(ext.get_storage(&key).data.is_none());
 
 		Ok(())
 	}
 
+	// c: is_cold (0 = hot, 1 = cold)
 	#[benchmark(skip_meta, pov_mode = Measured)]
-	fn seal_get_storage(n: Linear<0, { limits::STORAGE_BYTES }>) -> Result<(), BenchmarkError> {
+	fn seal_get_storage(
+		n: Linear<0, { limits::STORAGE_BYTES }>,
+		c: Linear<0, 1>,
+	) -> Result<(), BenchmarkError> {
 		let max_key_len = limits::STORAGE_KEY_BYTES;
 		let key = Key::try_from_var(vec![0u8; max_key_len as usize])
 			.map_err(|_| "Key has wrong length")?;
@@ -1486,6 +1511,15 @@ mod benchmarks {
 			.map_err(|_| "Failed to write to storage during setup.")?;
 
 		let out_ptr = max_key_len + 4;
+
+		// Whitelist key if c=0 (hot)
+		if c == 0 {
+			let mut full_key = info.child_trie_info().prefixed_storage_key().into_inner();
+			full_key.extend_from_slice(&key.hash());
+			frame_benchmarking::benchmarking::add_to_whitelist(full_key.into());
+			let _ = info.read(&key);
+		}
+
 		let result;
 		#[block]
 		{
@@ -1500,12 +1534,16 @@ mod benchmarks {
 		}
 
 		assert_ok!(result);
-		assert_eq!(&info.read(&key).unwrap(), &memory[out_ptr as usize..]);
+		assert_eq!(&info.read(&key).data.unwrap(), &memory[out_ptr as usize..]);
 		Ok(())
 	}
 
+	// c: is_cold (0 = hot, 1 = cold)
 	#[benchmark(skip_meta, pov_mode = Measured)]
-	fn contains_storage(n: Linear<0, { limits::STORAGE_BYTES }>) -> Result<(), BenchmarkError> {
+	fn contains_storage(
+		n: Linear<0, { limits::STORAGE_BYTES }>,
+		c: Linear<0, 1>,
+	) -> Result<(), BenchmarkError> {
 		let max_key_len = limits::STORAGE_KEY_BYTES;
 		let key = Key::try_from_var(vec![0u8; max_key_len as usize])
 			.map_err(|_| "Key has wrong length")?;
@@ -1517,6 +1555,15 @@ mod benchmarks {
 		.abi_encode();
 
 		let mut call_setup = CallSetup::<T>::default();
+
+		// Whitelist key if c=0 (hot)
+		if c == 0 {
+			let info = call_setup.contract().info()?;
+			let mut full_key = info.child_trie_info().prefixed_storage_key().into_inner();
+			full_key.extend_from_slice(&key.hash());
+			frame_benchmarking::benchmarking::add_to_whitelist(full_key.into());
+		}
+
 		let (mut ext, _) = call_setup.ext();
 		ext.set_storage(&key, Some(vec![42u8; max_key_len as usize]), false)
 			.map_err(|_| "Failed to write to storage during setup.")?;
@@ -1527,17 +1574,21 @@ mod benchmarks {
 			result = run_builtin_precompile(
 				&mut ext,
 				H160(BenchmarkStorage::<T>::MATCHER.base_address()).as_fixed_bytes(),
-				input_bytes,
+				input_bytes.clone(),
 			);
 		}
 		assert_ok!(result);
-		assert!(ext.get_storage(&key).is_some());
+		assert!(ext.get_storage(&key).data.is_some());
 
 		Ok(())
 	}
 
+	// c: is_cold (0 = hot, 1 = cold)
 	#[benchmark(skip_meta, pov_mode = Measured)]
-	fn take_storage(n: Linear<0, { limits::STORAGE_BYTES }>) -> Result<(), BenchmarkError> {
+	fn take_storage(
+		n: Linear<0, { limits::STORAGE_BYTES }>,
+		c: Linear<0, 1>,
+	) -> Result<(), BenchmarkError> {
 		let max_key_len = limits::STORAGE_KEY_BYTES;
 		let key = Key::try_from_var(vec![3u8; max_key_len as usize])
 			.map_err(|_| "Key has wrong length")?;
@@ -1550,6 +1601,15 @@ mod benchmarks {
 		.abi_encode();
 
 		let mut call_setup = CallSetup::<T>::default();
+
+		// Whitelist key if c=0 (hot)
+		if c == 0 {
+			let info = call_setup.contract().info()?;
+			let mut full_key = info.child_trie_info().prefixed_storage_key().into_inner();
+			full_key.extend_from_slice(&key.hash());
+			frame_benchmarking::benchmarking::add_to_whitelist(full_key.into());
+		}
+
 		let (mut ext, _) = call_setup.ext();
 		ext.set_storage(&key, Some(vec![42u8; max_key_len as usize]), false)
 			.map_err(|_| "Failed to write to storage during setup.")?;
@@ -1560,11 +1620,11 @@ mod benchmarks {
 			result = run_builtin_precompile(
 				&mut ext,
 				H160(BenchmarkStorage::<T>::MATCHER.base_address()).as_fixed_bytes(),
-				input_bytes,
+				input_bytes.clone(),
 			);
 		}
 		assert_ok!(result);
-		assert!(ext.get_storage(&key).is_none());
+		assert!(ext.get_storage(&key).data.is_none());
 
 		Ok(())
 	}

substrate/frame/revive/src/exec.rs

@@ -527,9 +527,9 @@ pub trait PrecompileExt: sealing::Sealed {
 
 	/// Returns the storage entry of the executing account by the given `key`.
 	///
-	/// Returns `None` if the `key` wasn't previously set by `set_storage` or
-	/// was deleted.
-	fn get_storage(&mut self, key: &Key) -> Option<Vec<u8>>;
+	/// Returns a `StateLoad` containing the value (or `None` if the key wasn't previously set
+	/// or was deleted) and whether it was a cold or hot load.
+	fn get_storage(&mut self, key: &Key) -> sp_io::StateLoad<Option<Vec<u8>>>;
 
 	/// Returns `Some(len)` (in bytes) if a storage item exists at `key`.
 	///
@@ -544,7 +544,7 @@ pub trait PrecompileExt: sealing::Sealed {
 		key: &Key,
 		value: Option<Vec<u8>>,
 		take_old: bool,
-	) -> Result<WriteOutcome, DispatchError>;
+	) -> Result<sp_io::StateLoad<WriteOutcome>, DispatchError>;
 
 	/// Charges `diff` from the meter.
 	fn charge_storage(&mut self, diff: &Diff) -> DispatchResult;
@@ -2461,7 +2461,7 @@ where
 		frame.frame_meter.eth_gas_left().unwrap_or_default().saturated_into::<u64>()
 	}
 
-	fn get_storage(&mut self, key: &Key) -> Option<Vec<u8>> {
+	fn get_storage(&mut self, key: &Key) -> sp_io::StateLoad<Option<Vec<u8>>> {
 		assert!(self.has_contract_info());
 		self.top_frame_mut().contract_info().read(key)
 	}
@@ -2476,7 +2476,7 @@ where
 		key: &Key,
 		value: Option<Vec<u8>>,
 		take_old: bool,
-	) -> Result<WriteOutcome, DispatchError> {
+	) -> Result<sp_io::StateLoad<WriteOutcome>, DispatchError> {
 		assert!(self.has_contract_info());
 		let frame = self.top_frame_mut();
 		frame.contract_info.get(&frame.account_id).write(