move state to common

core/src/commands/create_snapshot.rs

@@ -22,9 +22,10 @@ use sp_runtime::traits::{Block as BlockT, NumberFor};
 use substrate_rpc_client::{ws_client, StateApi};
 
 use crate::{
-    build_executor,
-    common::shared_parameters,
-    state::{LiveState, RuntimeChecks, State},
+    common::{
+        shared_parameters,
+        state::{build_executor, LiveState, RuntimeChecks, State},
+    },
     SharedParams, LOG_TARGET,
 };
 

core/src/commands/execute_block.rs

@@ -26,9 +26,10 @@ use sp_runtime::{
 use substrate_rpc_client::{ws_client, ChainApi};
 
 use crate::{
-    build_executor, full_extensions, rpc_err_handler,
-    state::{LiveState, RuntimeChecks, State},
-    state_machine_call_with_proof, SharedParams, LOG_TARGET,
+    common::state::{
+        build_executor, state_machine_call_with_proof, LiveState, RuntimeChecks, State,
+    },
+    full_extensions, rpc_err_handler, SharedParams, LOG_TARGET,
 };
 
 /// Configurations for [`run`].

core/src/commands/fast_forward.rs

@@ -30,13 +30,15 @@ use sp_runtime::{
 use sp_state_machine::TestExternalities;
 
 use crate::{
-    build_executor, full_extensions,
+    common::state::{
+        build_executor, state_machine_call, state_machine_call_with_proof, RuntimeChecks, State,
+    },
+    full_extensions,
     inherents::{
         pre_apply::pre_apply_inherents,
         providers::{InherentProvider, ProviderVariant},
     },
-    state::{RuntimeChecks, State},
-    state_machine_call, state_machine_call_with_proof, BlockT, SharedParams,
+    BlockT, SharedParams,
 };
 
 /// Configuration for [`run`].

core/src/commands/follow_chain.rs

@@ -28,11 +28,11 @@ use sp_runtime::{
 use substrate_rpc_client::{ws_client, ChainApi, FinalizedHeaders, Subscription, WsClient};
 
 use crate::{
-    build_executor,
-    common::parse,
-    full_extensions, rpc_err_handler,
-    state::{LiveState, RuntimeChecks, State},
-    state_machine_call_with_proof, SharedParams, LOG_TARGET,
+    common::{
+        parse,
+        state::{build_executor, state_machine_call_with_proof, LiveState, RuntimeChecks, State},
+    },
+    full_extensions, rpc_err_handler, SharedParams, LOG_TARGET,
 };
 
 const SUB: &str = "chain_subscribeFinalizedHeads";

core/src/commands/offchain_worker.rs

@@ -23,11 +23,11 @@ use sp_runtime::traits::{Block as BlockT, NumberFor};
 use substrate_rpc_client::{ws_client, ChainApi};
 
 use crate::{
-    build_executor,
-    common::parse,
-    full_extensions, rpc_err_handler,
-    state::{LiveState, RuntimeChecks, State},
-    state_machine_call, SharedParams, LOG_TARGET,
+    common::{
+        parse,
+        state::{build_executor, state_machine_call, LiveState, RuntimeChecks, State},
+    },
+    full_extensions, rpc_err_handler, SharedParams, LOG_TARGET,
 };
 
 /// Configuration for [`run`].

core/src/commands/on_runtime_upgrade.rs

@@ -27,10 +27,11 @@ use sp_runtime::traits::{Block as BlockT, HashingFor, NumberFor};
 use sp_state_machine::{CompactProof, OverlayedChanges, StorageProof};
 
 use crate::{
-    build_executor,
-    common::misc_logging,
-    state::{RuntimeChecks, State},
-    state_machine_call_with_proof, RefTimeInfo, SharedParams, LOG_TARGET,
+    common::{
+        misc_logging,
+        state::{build_executor, state_machine_call_with_proof, RuntimeChecks, State},
+    },
+    RefTimeInfo, SharedParams, LOG_TARGET,
 };
 
 /// Configuration for [`run`].

core/src/common/mod.rs

@@ -1,3 +1,4 @@
 pub mod misc_logging;
 pub mod parse;
 pub mod shared_parameters;
+pub mod state;

core/src/state.rs → core/src/common/state.rs

@@ -21,23 +21,28 @@ use frame_remote_externalities::{
     Builder, Mode, OfflineConfig, OnlineConfig, RemoteExternalities, SnapshotConfig,
 };
 use parity_scale_codec::Decode;
-use sc_cli::RuntimeVersion;
-use sc_executor::{sp_wasm_interface::HostFunctions, WasmExecutor};
+use sc_cli::{execution_method_from_cli, RuntimeVersion};
+use sc_executor::{
+    sp_wasm_interface::HostFunctions, HeapAllocStrategy, WasmExecutor, DEFAULT_HEAP_ALLOC_STRATEGY,
+};
+use sp_api::{CallContext, StorageProof};
 use sp_core::{
     hexdisplay::HexDisplay, storage::well_known_keys, traits::ReadRuntimeVersion, twox_128, Hasher,
 };
+use sp_externalities::Extensions;
 use sp_runtime::{
-    traits::{BlakeTwo256, Block as BlockT, Header as HeaderT},
+    traits::{BlakeTwo256, Block as BlockT, HashingFor, Header as HeaderT},
     DeserializeOwned,
 };
+use sp_state_machine::{OverlayedChanges, StateMachine, TestExternalities, TrieBackendBuilder};
 use substrate_rpc_client::{ws_client, ChainApi};
 
 use crate::{
     common::{
         parse,
         shared_parameters::{Runtime, SharedParams},
     },
-    ensure_try_runtime, hash_of, rpc_err_handler, LOG_TARGET,
+    hash_of, rpc_err_handler, LOG_TARGET,
 };
 
 /// A `Live` variant for [`State`]
@@ -315,3 +320,156 @@ impl State {
         Ok(ext)
     }
 }
+
+/// Build wasm executor by default config.
+pub(crate) fn build_executor<H: HostFunctions>(shared: &SharedParams) -> WasmExecutor<H> {
+    let heap_pages =
+        shared
+            .heap_pages
+            .map_or(DEFAULT_HEAP_ALLOC_STRATEGY, |p| HeapAllocStrategy::Static {
+                extra_pages: p as _,
+            });
+
+    WasmExecutor::builder()
+        .with_execution_method(execution_method_from_cli(
+            shared.wasm_method,
+            shared.wasmtime_instantiation_strategy,
+        ))
+        .with_onchain_heap_alloc_strategy(heap_pages)
+        .with_offchain_heap_alloc_strategy(heap_pages)
+        // There is not that much we can do if someone is using unknown host functions.
+        // They would need to fork the `cli` to add their custom host functions.
+        .with_allow_missing_host_functions(true)
+        .build()
+}
+
+/// Ensure that the given `ext` is compiled with `try-runtime`
+fn ensure_try_runtime<Block: BlockT, HostFns: HostFunctions>(
+    executor: &WasmExecutor<HostFns>,
+    ext: &mut TestExternalities<HashingFor<Block>>,
+) -> bool {
+    use sp_api::RuntimeApiInfo;
+    let final_code = ext
+        .execute_with(|| sp_io::storage::get(well_known_keys::CODE))
+        .expect("':CODE:' is always downloaded in try-runtime-cli; qed");
+    let final_version = <RuntimeVersion as Decode>::decode(
+        &mut &*executor
+            .read_runtime_version(&final_code, &mut ext.ext())
+            .unwrap(),
+    )
+    .unwrap();
+    final_version
+        .api_version(&<dyn frame_try_runtime::TryRuntime<Block>>::ID)
+        .is_some()
+}
+
+/// Execute the given `method` and `data` on top of `ext`, returning the results (encoded) and the
+/// state `changes`.
+pub(crate) fn state_machine_call<Block: BlockT, HostFns: HostFunctions>(
+    ext: &TestExternalities<HashingFor<Block>>,
+    executor: &WasmExecutor<HostFns>,
+    method: &'static str,
+    data: &[u8],
+    mut extensions: Extensions,
+) -> sc_cli::Result<(OverlayedChanges<HashingFor<Block>>, Vec<u8>)> {
+    let mut changes = Default::default();
+    let encoded_result = StateMachine::new(
+        &ext.backend,
+        &mut changes,
+        executor,
+        method,
+        data,
+        &mut extensions,
+        &sp_state_machine::backend::BackendRuntimeCode::new(&ext.backend).runtime_code()?,
+        CallContext::Offchain,
+    )
+    .execute()
+    .map_err(|e| format!("failed to execute '{}': {}", method, e))
+    .map_err::<sc_cli::Error, _>(Into::into)?;
+
+    Ok((changes, encoded_result))
+}
+
+/// Same as [`state_machine_call`], but it also computes and returns the storage proof and ref time
+/// information.
+///
+/// Make sure [`LOG_TARGET`] is enabled in logging.
+pub(crate) fn state_machine_call_with_proof<Block: BlockT, HostFns: HostFunctions>(
+    ext: &TestExternalities<HashingFor<Block>>,
+    storage_overlay: &mut OverlayedChanges<HashingFor<Block>>,
+    executor: &WasmExecutor<HostFns>,
+    method: &'static str,
+    data: &[u8],
+    mut extensions: Extensions,
+    maybe_export_proof: Option<PathBuf>,
+) -> sc_cli::Result<(StorageProof, Vec<u8>)> {
+    let runtime_code_backend = sp_state_machine::backend::BackendRuntimeCode::new(&ext.backend);
+    let proving_backend = TrieBackendBuilder::wrap(&ext.backend)
+        .with_recorder(Default::default())
+        .build();
+    let runtime_code = runtime_code_backend.runtime_code()?;
+
+    let encoded_result = StateMachine::new(
+        &proving_backend,
+        storage_overlay,
+        executor,
+        method,
+        data,
+        &mut extensions,
+        &runtime_code,
+        CallContext::Offchain,
+    )
+    .execute()
+    .map_err(|e| format!("failed to execute {}: {}", method, e))
+    .map_err::<sc_cli::Error, _>(Into::into)?;
+
+    let proof = proving_backend
+        .extract_proof()
+        .expect("A recorder was set and thus, a storage proof can be extracted; qed");
+
+    if let Some(path) = maybe_export_proof {
+        let mut file = std::fs::File::create(&path).map_err(|e| {
+            log::error!(
+                target: LOG_TARGET,
+                "Failed to create file {}: {:?}",
+                path.to_string_lossy(),
+                e
+            );
+            e
+        })?;
+
+        log::info!(target: LOG_TARGET, "Writing storage proof to {}", path.to_string_lossy());
+
+        use std::io::Write as _;
+        file.write_all(storage_proof_to_raw_json(&proof).as_bytes())
+            .map_err(|e| {
+                log::error!(
+                    target: LOG_TARGET,
+                    "Failed to write storage proof to {}: {:?}",
+                    path.to_string_lossy(),
+                    e
+                );
+                e
+            })?;
+    }
+
+    Ok((proof, encoded_result))
+}
+
+/// Converts a [`sp_state_machine::StorageProof`] into a JSON string.
+fn storage_proof_to_raw_json(storage_proof: &sp_state_machine::StorageProof) -> String {
+    serde_json::Value::Object(
+        storage_proof
+            .to_memory_db::<sp_runtime::traits::BlakeTwo256>()
+            .drain()
+            .iter()
+            .map(|(key, (value, _n))| {
+                (
+                    format!("0x{}", hex::encode(key.as_bytes())),
+                    serde_json::Value::String(format!("0x{}", hex::encode(value))),
+                )
+            })
+            .collect(),
+    )
+    .to_string()
+}