From 1ed54ff0dc425b22027ce138d46b9e5ce324ac2b Mon Sep 17 00:00:00 2001
From: librelois <c@elo.tf>
Date: Wed, 27 Sep 2023 15:01:10 +0200
Subject: [PATCH] improve precompile-wasm subcommand: load the bytecode from
 chain spec

---
 Cargo.lock                                    |  1 +
 .../bin/node-template/node/src/command.rs     |  2 +-
 substrate/client/cli/Cargo.toml               |  1 +
 .../cli/src/commands/precompile_wasm_cmd.rs   | 63 +++++++++++++------
 4 files changed, 47 insertions(+), 20 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 8064516baefa..44cc8fe9d37f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -14673,6 +14673,7 @@ dependencies = [
  "sp-panic-handler",
  "sp-runtime",
  "sp-state-machine",
+ "sp-storage",
  "sp-tracing",
  "sp-version",
  "tempfile",
diff --git a/substrate/bin/node-template/node/src/command.rs b/substrate/bin/node-template/node/src/command.rs
index 097f7b23fa6b..b722e7e21097 100644
--- a/substrate/bin/node-template/node/src/command.rs
+++ b/substrate/bin/node-template/node/src/command.rs
@@ -183,7 +183,7 @@ pub fn run() -> sc_cli::Result<()> {
 			runner.async_run(|config| {
 				let PartialComponents { task_manager, backend, .. } =
 					service::new_partial(&config)?;
-				Ok((cmd.run(backend), task_manager))
+				Ok((cmd.run(backend, config.chain_spec), task_manager))
 			})
 		},
 		None => {
diff --git a/substrate/client/cli/Cargo.toml b/substrate/client/cli/Cargo.toml
index 10e21404b534..8e95714bbf51 100644
--- a/substrate/client/cli/Cargo.toml
+++ b/substrate/client/cli/Cargo.toml
@@ -46,6 +46,7 @@ sp-keystore = { path = "../../primitives/keystore" }
 sp-panic-handler = { path = "../../primitives/panic-handler" }
 sp-runtime = { path = "../../primitives/runtime" }
 sp-state-machine = { path = "../../primitives/state-machine" }
+sp-storage = { path = "../../primitives/storage" }
 sp-version = { path = "../../primitives/version" }
 
 [dev-dependencies]
diff --git a/substrate/client/cli/src/commands/precompile_wasm_cmd.rs b/substrate/client/cli/src/commands/precompile_wasm_cmd.rs
index 7e5bedaad115..f0366f35f572 100644
--- a/substrate/client/cli/src/commands/precompile_wasm_cmd.rs
+++ b/substrate/client/cli/src/commands/precompile_wasm_cmd.rs
@@ -28,10 +28,12 @@ use crate::{
 
 use clap::Parser;
 use sc_client_api::{Backend, HeaderBackend};
+use sp_core::traits::RuntimeCode;
 use sc_executor::{
-	precompile_and_serialize_versioned_wasm_runtime, HeapAllocStrategy, DEFAULT_HEAP_ALLOC_STRATEGY,
+	precompile_and_serialize_versioned_wasm_runtime, HeapAllocStrategy, DEFAULT_HEAP_ALLOC_PAGES,
 };
 use sp_runtime::traits::Block as BlockT;
+use sc_service::ChainSpec;
 use sp_state_machine::backend::BackendRuntimeCode;
 use std::{fmt::Debug, path::PathBuf, sync::Arc};
 
@@ -79,28 +81,51 @@ pub struct PrecompileWasmCmd {
 
 impl PrecompileWasmCmd {
 	/// Run the precompile-wasm command
-	pub async fn run<B, BA>(&self, backend: Arc<BA>) -> error::Result<()>
+	pub async fn run<B, BA>(&self, backend: Arc<BA>, spec: Box<dyn ChainSpec>,) -> error::Result<()>
 	where
 		B: BlockT,
 		BA: Backend<B>,
 	{
-		let state = backend.state_at(backend.blockchain().info().finalized_hash)?;
-
-		let heap_pages = self.default_heap_pages.map_or(DEFAULT_HEAP_ALLOC_STRATEGY, |h| {
-			HeapAllocStrategy::Static { extra_pages: h as _ }
-		});
-
-		precompile_and_serialize_versioned_wasm_runtime(
-			true,
-			heap_pages,
-			&BackendRuntimeCode::new(&state).runtime_code()?,
-			execution_method_from_cli(
-				WasmExecutionMethod::Compiled,
-				self.wasmtime_instantiation_strategy,
-			),
-			&self.output_dir,
-		)
-		.map_err(|e| Error::Application(Box::new(e)))?;
+		let heap_pages = self.default_heap_pages.unwrap_or(DEFAULT_HEAP_ALLOC_PAGES);
+
+		let blockchain_info = backend.blockchain().info();
+
+		if backend.have_state_at(blockchain_info.finalized_hash, blockchain_info.finalized_number) {
+			let state = backend.state_at(backend.blockchain().info().finalized_hash)?;
+	
+			precompile_and_serialize_versioned_wasm_runtime(
+				true,
+				HeapAllocStrategy::Static { extra_pages: heap_pages },
+				&BackendRuntimeCode::new(&state).runtime_code()?,
+				execution_method_from_cli(
+					WasmExecutionMethod::Compiled,
+					self.wasmtime_instantiation_strategy,
+				),
+				&self.output_dir,
+			)
+			.map_err(|e| Error::Application(Box::new(e)))?;
+		} else {
+			let storage = spec.as_storage_builder().build_storage()?;
+			if let Some(wasm_bytecode) = storage.top.get(sp_storage::well_known_keys::CODE) {
+				let runtime_code = RuntimeCode {
+					code_fetcher: &sp_core::traits::WrappedRuntimeCode(wasm_bytecode.as_slice().into()),
+					hash: sp_core::blake2_256(&wasm_bytecode).to_vec(),
+					heap_pages: Some(heap_pages as u64),
+				};
+				precompile_and_serialize_versioned_wasm_runtime(
+					true,
+					HeapAllocStrategy::Static { extra_pages: heap_pages },
+					&runtime_code,
+					execution_method_from_cli(
+						WasmExecutionMethod::Compiled,
+						self.wasmtime_instantiation_strategy,
+					),
+					&self.output_dir,
+				)
+				.map_err(|e| Error::Application(Box::new(e)))?;
+			}
+		}
+
 
 		Ok(())
 	}