BlockId removal: refactor: Backend::state_at (paritytech#12488)

* Minor naming improved

* BlockId removal refactor: Backend::state_at

* formatting

bin/node/bench/src/import.rs

@@ -34,7 +34,7 @@ use std::borrow::Cow;
 
 use node_primitives::Block;
 use node_testing::bench::{BenchDb, BlockType, DatabaseType, KeyTypes, Profile};
-use sc_client_api::backend::Backend;
+use sc_client_api::{backend::Backend, HeaderBackend};
 use sp_runtime::generic::BlockId;
 use sp_state_machine::InspectState;
 
@@ -127,10 +127,15 @@ impl core::Benchmark for ImportBenchmark {
 		context.import_block(self.block.clone());
 		let elapsed = start.elapsed();
 
+		let hash = context
+			.client
+			.expect_block_hash_from_id(&BlockId::number(1))
+			.expect("Block 1 was imported; qed");
+
 		// Sanity checks.
 		context
 			.client
-			.state_at(&BlockId::number(1))
+			.state_at(&hash)
 			.expect("state_at failed for block#1")
 			.inspect_state(|| {
 				match self.block_type {

client/api/src/backend.rs

@@ -505,11 +505,11 @@ pub trait Backend<Block: BlockT>: AuxStore + Send + Sync {
 
 	/// Returns true if state for given block is available.
 	fn have_state_at(&self, hash: &Block::Hash, _number: NumberFor<Block>) -> bool {
-		self.state_at(BlockId::Hash(*hash)).is_ok()
+		self.state_at(hash).is_ok()
 	}
 
 	/// Returns state backend with post-state of given block.
-	fn state_at(&self, block: BlockId<Block>) -> sp_blockchain::Result<Self::State>;
+	fn state_at(&self, hash: &Block::Hash) -> sp_blockchain::Result<Self::State>;
 
 	/// Attempts to revert the chain by `n` blocks. If `revert_finalized` is set it will attempt to
 	/// revert past any finalized block, this is unsafe and can potentially leave the node in an

client/api/src/in_mem.rs

@@ -686,7 +686,7 @@ where
 	type OffchainStorage = OffchainStorage;
 
 	fn begin_operation(&self) -> sp_blockchain::Result<Self::BlockImportOperation> {
-		let old_state = self.state_at(BlockId::Hash(Default::default()))?;
+		let old_state = self.state_at(&Default::default())?;
 		Ok(BlockImportOperation {
 			pending_block: None,
 			old_state,
@@ -702,7 +702,8 @@ where
 		operation: &mut Self::BlockImportOperation,
 		block: BlockId<Block>,
 	) -> sp_blockchain::Result<()> {
-		operation.old_state = self.state_at(block)?;
+		let hash = self.blockchain.expect_block_hash_from_id(&block)?;
+		operation.old_state = self.state_at(&hash)?;
 		Ok(())
 	}
 
@@ -768,16 +769,16 @@ where
 		None
 	}
 
-	fn state_at(&self, block: BlockId<Block>) -> sp_blockchain::Result<Self::State> {
-		match block {
-			BlockId::Hash(h) if h == Default::default() => return Ok(Self::State::default()),
-			_ => {},
+	fn state_at(&self, hash: &Block::Hash) -> sp_blockchain::Result<Self::State> {
+		if *hash == Default::default() {
+			return Ok(Self::State::default())
 		}
 
-		self.blockchain
-			.id(block)
-			.and_then(|id| self.states.read().get(&id).cloned())
-			.ok_or_else(|| sp_blockchain::Error::UnknownBlock(format!("{}", block)))
+		self.states
+			.read()
+			.get(hash)
+			.cloned()
+			.ok_or_else(|| sp_blockchain::Error::UnknownBlock(format!("{}", hash)))
 	}
 
 	fn revert(

client/basic-authorship/src/basic_authorship.rs

@@ -736,7 +736,7 @@ mod tests {
 		let api = client.runtime_api();
 		api.execute_block(&block_id, proposal.block).unwrap();
 
-		let state = backend.state_at(block_id).unwrap();
+		let state = backend.state_at(&genesis_hash).unwrap();
 
 		let storage_changes = api.into_storage_changes(&state, genesis_hash).unwrap();
 

client/block-builder/src/lib.rs

@@ -258,12 +258,11 @@ where
 
 		let proof = self.api.extract_proof();
 
-		let state = self.backend.state_at(self.block_id)?;
-		let parent_hash = self.parent_hash;
+		let state = self.backend.state_at(&self.parent_hash)?;
 
 		let storage_changes = self
 			.api
-			.into_storage_changes(&state, parent_hash)
+			.into_storage_changes(&state, self.parent_hash)
 			.map_err(sp_blockchain::Error::StorageChanges)?;
 
 		Ok(BuiltBlock {

client/db/benches/state_access.rs

@@ -84,7 +84,7 @@ fn insert_blocks(db: &Backend<Block>, storage: Vec<(Vec<u8>, Vec<u8>)>) -> H256
 			.map(|(k, v)| (k.clone(), Some(v.clone())))
 			.collect::<Vec<_>>();
 
-		let (state_root, tx) = db.state_at(BlockId::Hash(parent_hash)).unwrap().storage_root(
+		let (state_root, tx) = db.state_at(&parent_hash).unwrap().storage_root(
 			changes.iter().map(|(k, v)| (k.as_slice(), v.as_deref())),
 			StateVersion::V1,
 		);
@@ -176,7 +176,7 @@ fn state_access_benchmarks(c: &mut Criterion) {
 
 		group.bench_function(desc, |b| {
 			b.iter_batched(
-				|| backend.state_at(BlockId::Hash(block_hash)).expect("Creates state"),
+				|| backend.state_at(&block_hash).expect("Creates state"),
 				|state| {
 					for key in keys.iter().cycle().take(keys.len() * multiplier) {
 						let _ = state.storage(&key).expect("Doesn't fail").unwrap();
@@ -214,7 +214,7 @@ fn state_access_benchmarks(c: &mut Criterion) {
 
 		group.bench_function(desc, |b| {
 			b.iter_batched(
-				|| backend.state_at(BlockId::Hash(block_hash)).expect("Creates state"),
+				|| backend.state_at(&block_hash).expect("Creates state"),
 				|state| {
 					for key in keys.iter().take(1).cycle().take(multiplier) {
 						let _ = state.storage(&key).expect("Doesn't fail").unwrap();
@@ -252,7 +252,7 @@ fn state_access_benchmarks(c: &mut Criterion) {
 
 		group.bench_function(desc, |b| {
 			b.iter_batched(
-				|| backend.state_at(BlockId::Hash(block_hash)).expect("Creates state"),
+				|| backend.state_at(&block_hash).expect("Creates state"),
 				|state| {
 					for key in keys.iter().take(1).cycle().take(multiplier) {
 						let _ = state.storage_hash(&key).expect("Doesn't fail").unwrap();
@@ -290,7 +290,7 @@ fn state_access_benchmarks(c: &mut Criterion) {
 
 		group.bench_function(desc, |b| {
 			b.iter_batched(
-				|| backend.state_at(BlockId::Hash(block_hash)).expect("Creates state"),
+				|| backend.state_at(&block_hash).expect("Creates state"),
 				|state| {
 					let _ = state
 						.storage_hash(sp_core::storage::well_known_keys::CODE)

client/db/src/lib.rs

@@ -1963,10 +1963,11 @@ impl<Block: BlockT> sc_client_api::backend::Backend<Block> for Backend<Block> {
 		operation: &mut Self::BlockImportOperation,
 		block: BlockId<Block>,
 	) -> ClientResult<()> {
+		let hash = self.blockchain.expect_block_hash_from_id(&block)?;
 		if block.is_pre_genesis() {
 			operation.old_state = self.empty_state()?;
 		} else {
-			operation.old_state = self.state_at(block)?;
+			operation.old_state = self.state_at(&hash)?;
 		}
 
 		operation.commit_state = true;
@@ -2302,15 +2303,8 @@ impl<Block: BlockT> sc_client_api::backend::Backend<Block> for Backend<Block> {
 		&self.blockchain
 	}
 
-	fn state_at(&self, block: BlockId<Block>) -> ClientResult<Self::State> {
-		use sc_client_api::blockchain::HeaderBackend as BcHeaderBackend;
-
-		let is_genesis = match &block {
-			BlockId::Number(n) if n.is_zero() => true,
-			BlockId::Hash(h) if h == &self.blockchain.meta.read().genesis_hash => true,
-			_ => false,
-		};
-		if is_genesis {
+	fn state_at(&self, hash: &Block::Hash) -> ClientResult<Self::State> {
+		if hash == &self.blockchain.meta.read().genesis_hash {
 			if let Some(genesis_state) = &*self.genesis_state.read() {
 				let root = genesis_state.root;
 				let db_state = DbStateBuilder::<Block>::new(genesis_state.clone(), root)
@@ -2322,35 +2316,28 @@ impl<Block: BlockT> sc_client_api::backend::Backend<Block> for Backend<Block> {
 			}
 		}
 
-		let hash = match block {
-			BlockId::Hash(h) => h,
-			BlockId::Number(n) => self.blockchain.hash(n)?.ok_or_else(|| {
-				sp_blockchain::Error::UnknownBlock(format!("Unknown block number {}", n))
-			})?,
-		};
-
-		match self.blockchain.header_metadata(hash) {
+		match self.blockchain.header_metadata(*hash) {
 			Ok(ref hdr) => {
 				let hint = || {
 					sc_state_db::NodeDb::get(self.storage.as_ref(), hdr.state_root.as_ref())
 						.unwrap_or(None)
 						.is_some()
 				};
 				if let Ok(()) =
-					self.storage.state_db.pin(&hash, hdr.number.saturated_into::<u64>(), hint)
+					self.storage.state_db.pin(hash, hdr.number.saturated_into::<u64>(), hint)
 				{
 					let root = hdr.state_root;
 					let db_state = DbStateBuilder::<Block>::new(self.storage.clone(), root)
 						.with_optional_cache(
 							self.shared_trie_cache.as_ref().map(|c| c.local_cache()),
 						)
 						.build();
-					let state = RefTrackingState::new(db_state, self.storage.clone(), Some(hash));
-					Ok(RecordStatsState::new(state, Some(hash), self.state_usage.clone()))
+					let state = RefTrackingState::new(db_state, self.storage.clone(), Some(*hash));
+					Ok(RecordStatsState::new(state, Some(*hash), self.state_usage.clone()))
 				} else {
 					Err(sp_blockchain::Error::UnknownBlock(format!(
 						"State already discarded for {:?}",
-						block
+						hash
 					)))
 				}
 			},
@@ -2588,7 +2575,7 @@ pub(crate) mod tests {
 
 			db.commit_operation(op).unwrap();
 
-			let state = db.state_at(BlockId::Number(0)).unwrap();
+			let state = db.state_at(&hash).unwrap();
 
 			assert_eq!(state.storage(&[1, 3, 5]).unwrap(), Some(vec![2, 4, 6]));
 			assert_eq!(state.storage(&[1, 2, 3]).unwrap(), Some(vec![9, 9, 9]));
@@ -2623,7 +2610,8 @@ pub(crate) mod tests {
 
 			db.commit_operation(op).unwrap();
 
-			let state = db.state_at(BlockId::Number(1)).unwrap();
+			let hash = db.blockchain().expect_block_hash_from_id(&BlockId::Number(1)).unwrap();
+			let state = db.state_at(&hash).unwrap();
 
 			assert_eq!(state.storage(&[1, 3, 5]).unwrap(), None);
 			assert_eq!(state.storage(&[1, 2, 3]).unwrap(), Some(vec![9, 9, 9]));
@@ -3139,11 +3127,7 @@ pub(crate) mod tests {
 			hash
 		};
 
-		let block0_hash = backend
-			.state_at(BlockId::Hash(hash0))
-			.unwrap()
-			.storage_hash(&b"test"[..])
-			.unwrap();
+		let block0_hash = backend.state_at(&hash0).unwrap().storage_hash(&b"test"[..]).unwrap();
 
 		let hash1 = {
 			let mut op = backend.begin_operation().unwrap();
@@ -3182,11 +3166,7 @@ pub(crate) mod tests {
 			backend.commit_operation(op).unwrap();
 		}
 
-		let block1_hash = backend
-			.state_at(BlockId::Hash(hash1))
-			.unwrap()
-			.storage_hash(&b"test"[..])
-			.unwrap();
+		let block1_hash = backend.state_at(&hash1).unwrap().storage_hash(&b"test"[..]).unwrap();
 
 		assert_ne!(block0_hash, block1_hash);
 	}

client/service/src/client/call_executor.rs

@@ -147,17 +147,14 @@ where
 		extensions: Option<Extensions>,
 	) -> sp_blockchain::Result<Vec<u8>> {
 		let mut changes = OverlayedChanges::default();
-		let state = self.backend.state_at(*at)?;
+		let at_hash = self.backend.blockchain().expect_block_hash_from_id(at)?;
+		let state = self.backend.state_at(&at_hash)?;
 		let state_runtime_code = sp_state_machine::backend::BackendRuntimeCode::new(&state);
 		let runtime_code =
 			state_runtime_code.runtime_code().map_err(sp_blockchain::Error::RuntimeCode)?;
 
 		let runtime_code = self.check_override(runtime_code, at)?;
 
-		let at_hash = self.backend.blockchain().block_hash_from_id(at)?.ok_or_else(|| {
-			sp_blockchain::Error::UnknownBlock(format!("Could not find block hash for {:?}", at))
-		})?;
-
 		let mut sm = StateMachine::new(
 			&state,
 			&mut changes,
@@ -195,14 +192,11 @@ where
 	{
 		let mut storage_transaction_cache = storage_transaction_cache.map(|c| c.borrow_mut());
 
-		let state = self.backend.state_at(*at)?;
+		let at_hash = self.backend.blockchain().expect_block_hash_from_id(at)?;
+		let state = self.backend.state_at(&at_hash)?;
 
 		let changes = &mut *changes.borrow_mut();
 
-		let at_hash = self.backend.blockchain().block_hash_from_id(at)?.ok_or_else(|| {
-			sp_blockchain::Error::UnknownBlock(format!("Could not find block hash for {:?}", at))
-		})?;
-
 		// It is important to extract the runtime code here before we create the proof
 		// recorder to not record it. We also need to fetch the runtime code from `state` to
 		// make sure we use the caching layers.
@@ -255,7 +249,9 @@ where
 
 	fn runtime_version(&self, id: &BlockId<Block>) -> sp_blockchain::Result<RuntimeVersion> {
 		let mut overlay = OverlayedChanges::default();
-		let state = self.backend.state_at(*id)?;
+
+		let at_hash = self.backend.blockchain().expect_block_hash_from_id(id)?;
+		let state = self.backend.state_at(&at_hash)?;
 		let mut cache = StorageTransactionCache::<Block, B::State>::default();
 		let mut ext = Ext::new(&mut overlay, &mut cache, &state, None);
 		let state_runtime_code = sp_state_machine::backend::BackendRuntimeCode::new(&state);
@@ -272,7 +268,8 @@ where
 		method: &str,
 		call_data: &[u8],
 	) -> sp_blockchain::Result<(Vec<u8>, StorageProof)> {
-		let state = self.backend.state_at(*at)?;
+		let at_hash = self.backend.blockchain().expect_block_hash_from_id(at)?;
+		let state = self.backend.state_at(&at_hash)?;
 
 		let trie_backend = state.as_trie_backend();