try-runtime - support all state versions (paritytech#12089)

* Add argument

* Apply

* More verbose parsing

* fmt

* fmt again

* Invalid state version

* reuse parsing

* type mismatch

utils/frame/try-runtime/cli/src/commands/execute_block.rs

@@ -139,7 +139,8 @@ where
 			.state
 			.builder::<Block>()?
 			// make sure the state is being build with the parent hash, if it is online.
-			.overwrite_online_at(parent_hash.to_owned());
+			.overwrite_online_at(parent_hash.to_owned())
+			.state_version(shared.state_version);
 
 		let builder = if command.overwrite_wasm_code {
 			log::info!(

utils/frame/try-runtime/cli/src/commands/follow_chain.rs

@@ -102,11 +102,13 @@ where
 
 		// create an ext at the state of this block, whatever is the first subscription event.
 		if maybe_state_ext.is_none() {
-			let builder = Builder::<Block>::new().mode(Mode::Online(OnlineConfig {
-				transport: command.uri.clone().into(),
-				at: Some(*header.parent_hash()),
-				..Default::default()
-			}));
+			let builder = Builder::<Block>::new()
+				.mode(Mode::Online(OnlineConfig {
+					transport: command.uri.clone().into(),
+					at: Some(*header.parent_hash()),
+					..Default::default()
+				}))
+				.state_version(shared.state_version);
 
 			let new_ext = builder
 				.inject_hashed_key_value(&[(code_key.clone(), code.clone())])

utils/frame/try-runtime/cli/src/commands/offchain_worker.rs

@@ -128,7 +128,7 @@ where
 	);
 
 	let ext = {
-		let builder = command.state.builder::<Block>()?;
+		let builder = command.state.builder::<Block>()?.state_version(shared.state_version);
 
 		let builder = if command.overwrite_wasm_code {
 			log::info!(

utils/frame/try-runtime/cli/src/commands/on_runtime_upgrade.rs

@@ -52,7 +52,7 @@ where
 	let execution = shared.execution;
 
 	let ext = {
-		let builder = command.state.builder::<Block>()?;
+		let builder = command.state.builder::<Block>()?.state_version(shared.state_version);
 		let (code_key, code) = extract_code(&config.chain_spec)?;
 		builder.inject_hashed_key_value(&[(code_key, code)]).build().await?
 	};

utils/frame/try-runtime/cli/src/lib.rs

@@ -294,6 +294,7 @@ use sp_runtime::{
 	DeserializeOwned,
 };
 use sp_state_machine::{OverlayedChanges, StateMachine, TrieBackendBuilder};
+use sp_version::StateVersion;
 use std::{fmt::Debug, path::PathBuf, str::FromStr};
 
 mod commands;
@@ -421,6 +422,10 @@ pub struct SharedParams {
 	/// When enabled, the spec name check will not panic, and instead only show a warning.
 	#[clap(long)]
 	pub no_spec_name_check: bool,
+
+	/// State version that is used by the chain.
+	#[clap(long, default_value = "1", parse(try_from_str = parse::state_version))]
+	pub state_version: StateVersion,
 }
 
 /// Our `try-runtime` command.

utils/frame/try-runtime/cli/src/parse.rs

@@ -17,6 +17,8 @@
 
 //! Utils for parsing user input
 
+use sp_version::StateVersion;
+
 pub(crate) fn hash(block_hash: &str) -> Result<String, String> {
 	let (block_hash, offset) = if let Some(block_hash) = block_hash.strip_prefix("0x") {
 		(block_hash, 2)
@@ -42,3 +44,10 @@ pub(crate) fn url(s: &str) -> Result<String, &'static str> {
 		Err("not a valid WS(S) url: must start with 'ws://' or 'wss://'")
 	}
 }
+
+pub(crate) fn state_version(s: &str) -> Result<StateVersion, &'static str> {
+	s.parse::<u8>()
+		.map_err(|_| ())
+		.and_then(StateVersion::try_from)
+		.map_err(|_| "Invalid state version.")
+}