feat: pop bench pallet logic implementation

Cargo.toml

@@ -49,8 +49,13 @@ subxt = "0.38.0"
 ink_env = "5.0.0"
 sp-core = "32.0.0"
 sp-weights = "31.0.0"
-scale-info = { version = "2.11.4", default-features = false, features = ["derive"] }
-scale-value = { version = "0.17.0", default-features = false, features = ["from-string", "parser-ss58"] }
+scale-info = { version = "2.11.4", default-features = false, features = [
+    "derive",
+] }
+scale-value = { version = "0.17.0", default-features = false, features = [
+    "from-string",
+    "parser-ss58",
+] }
 contract-build = "5.0.2"
 contract-extrinsics = "5.0.2"
 contract-transcode = "5.0.2"
@@ -68,6 +73,12 @@ serde = { version = "1.0", features = ["derive"] }
 zombienet-sdk = "0.2.18"
 git2_credentials = "0.13.0"
 
+# benchmarking
+cumulus-primitives-proof-size-hostfunction = "0.10.0"
+frame-benchmarking-cli = "46.0.0"
+sp-runtime = "40.0.1"
+sp-statement-store = "15.0.0"
+
 # pop-cli
 clap = { version = "4.5", features = ["derive"] }
 cliclack = "0.3.1"

crates/pop-cli/Cargo.toml

@@ -52,6 +52,12 @@ axum.workspace = true
 open.workspace = true
 tower-http = { workspace = true, features = ["fs", "cors"] }
 
+# benchmarking
+cumulus-primitives-proof-size-hostfunction.workspace = true
+frame-benchmarking-cli.workspace = true
+sp-runtime.workspace = true
+sp-statement-store.workspace = true
+
 [dev-dependencies]
 assert_cmd.workspace = true
 contract-extrinsics.workspace = true

crates/pop-cli/src/commands/bench/mod.rs

@@ -0,0 +1,161 @@
+// SPDX-License-Identifier: GPL-3.0
+
+use crate::{
+	cli::{self, traits::Cli},
+	common::prompt::display_message,
+};
+use clap::{Args, Subcommand};
+use frame_benchmarking_cli::PalletCmd;
+use pop_parachains::generate_benchmarks;
+
+/// Arguments for bencharmking a project.
+#[derive(Args)]
+#[command(args_conflicts_with_subcommands = true)]
+pub struct BenchmarkArgs {
+	#[command(subcommand)]
+	pub command: Command,
+}
+
+/// Benchmark a pallet or parachain.
+#[derive(Subcommand)]
+pub enum Command {
+	/// Benchmark the extrinsic weight of FRAME Pallets
+	#[clap(alias = "p")]
+	Pallet(PalletCmd),
+}
+
+impl Command {
+	/// Executes the command.
+	pub(crate) fn execute(args: BenchmarkArgs) -> anyhow::Result<()> {
+		let mut cli = cli::Cli;
+
+		let result = match args.command {
+			Command::Pallet(cmd) => Command::bechmark_pallet(cmd, &mut cli),
+		};
+		match result {
+			Ok(()) => display_message("Benchmark completed successfully!", true, &mut cli),
+			Err(e) => display_message(&e.to_string(), false, &mut cli),
+		}
+	}
+
+	fn bechmark_pallet(cmd: PalletCmd, cli: &mut impl Cli) -> anyhow::Result<()> {
+		cli.intro("Benchmarking your pallets")?;
+		cli.warning(
+			"NOTE: the `pop bench pallet` is not yet battle tested - double check the results.",
+		)?;
+
+		if let Some(spec) = cmd.shared_params.chain {
+			return display_message(
+				&format!(
+					"Chain specs are not supported. Please remove `--chain={spec}` \
+					       and use `--runtime=<PATH>` instead"
+				),
+				false,
+				cli,
+			);
+		}
+		cli.warning("NOTE: this may take some time...")?;
+
+		if let Err(e) = generate_benchmarks(&cmd) {
+			return display_message(&e.to_string(), false, cli);
+		}
+
+		if let Some(output_path) = cmd.output {
+			console::Term::stderr().clear_last_lines(1)?;
+			cli.info(format!("Weight file is generated to {}", output_path.as_path().display()))?;
+		}
+		Ok(())
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	use crate::cli::MockCli;
+	use clap::Parser;
+	use std::env;
+
+	#[test]
+	fn benchmark_pallet_works() -> anyhow::Result<()> {
+		let mut cli = MockCli::new()
+			.expect_intro("Benchmarking your pallets")
+			.expect_warning(
+				"NOTE: the `pop bench pallet` is not yet battle tested - double check the results.",
+			)
+			.expect_warning("NOTE: this may take some time...");
+
+		let cmd = PalletCmd::try_parse_from(&[
+			"",
+			"--runtime",
+			get_mock_runtime_path().to_str().unwrap(),
+			"--pallet",
+			"pallet_timestamp",
+			"--extrinsic",
+			"",
+		])?;
+		Command::bechmark_pallet(cmd, &mut cli)?;
+		cli.verify()?;
+		Ok(())
+	}
+
+	#[test]
+	fn benchmark_pallet_with_chainspec_fails() -> anyhow::Result<()> {
+		let spec = "path-to-chainspec";
+		let expected = format!(
+			"Chain specs are not supported. Please remove `--chain={spec}` \
+		          and use `--runtime=<PATH>` instead"
+		);
+		let mut cli = MockCli::new()
+			.expect_intro("Benchmarking your pallets")
+			.expect_warning(
+				"NOTE: the `pop bench pallet` is not yet battle tested - double check the results.",
+			)
+			.expect_outro_cancel(expected.clone());
+
+		let cmd = PalletCmd::try_parse_from(&[
+			"",
+			"--chain",
+			spec,
+			"--pallet",
+			"pallet_timestamp",
+			"--extrinsic",
+			"",
+		])?;
+
+		Command::bechmark_pallet(cmd, &mut cli)?;
+		cli.verify()?;
+		Ok(())
+	}
+
+	#[test]
+	fn benchmark_pallet_fails_with_error() -> anyhow::Result<()> {
+		let mut cli = MockCli::new()
+			.expect_intro("Benchmarking your pallets")
+			.expect_warning(
+				"NOTE: the `pop bench pallet` is not yet battle tested - double check the results.",
+			)
+			.expect_outro_cancel(format!(
+				"Failed to run benchmarking: Invalid input: No benchmarks found which match your input."
+			));
+
+		let cmd = PalletCmd::try_parse_from(&[
+			"",
+			"--runtime",
+			get_mock_runtime_path().to_str().unwrap(),
+			"--pallet",
+			"unknown-pallet-name",
+			"--extrinsic",
+			"",
+		])?;
+
+		Command::bechmark_pallet(cmd, &mut cli)?;
+		cli.verify()?;
+		Ok(())
+	}
+
+	// Construct the path to the mock runtime WASM file.
+	fn get_mock_runtime_path() -> std::path::PathBuf {
+		env::current_dir().unwrap().join("tests/files/base_parachain.wasm")
+	}
+}

crates/pop-cli/src/commands/call/chain.rs

@@ -4,7 +4,10 @@ use std::path::Path;
 
 use crate::{
 	cli::{self, traits::*},
-	common::wallet::{prompt_to_use_wallet, request_signature},
+	common::{
+		prompt::display_message,
+		wallet::{prompt_to_use_wallet, request_signature},
+	},
 };
 use anyhow::{anyhow, Result};
 use clap::Args;
@@ -499,16 +502,6 @@ async fn submit_extrinsic_with_wallet(
 	Ok(())
 }
 
-// Displays a message to the user, with formatting based on the success status.
-fn display_message(message: &str, success: bool, cli: &mut impl Cli) -> Result<()> {
-	if success {
-		cli.outro(message)?;
-	} else {
-		cli.outro_cancel(message)?;
-	}
-	Ok(())
-}
-
 // Prompts the user for some predefined actions.
 fn prompt_predefined_actions(pallets: &[Pallet], cli: &mut impl Cli) -> Result<Option<Action>> {
 	let mut predefined_action = cli.select("What would you like to do?");
@@ -995,16 +988,6 @@ mod tests {
 		Ok(())
 	}
 
-	#[test]
-	fn display_message_works() -> Result<()> {
-		let mut cli = MockCli::new().expect_outro(&"Call completed successfully!");
-		display_message("Call completed successfully!", true, &mut cli)?;
-		cli.verify()?;
-		let mut cli = MockCli::new().expect_outro_cancel("Call failed.");
-		display_message("Call failed.", false, &mut cli)?;
-		cli.verify()
-	}
-
 	#[tokio::test]
 	async fn prompt_predefined_actions_works() -> Result<()> {
 		let client = set_up_client(POP_NETWORK_TESTNET_URL).await?;

crates/pop-cli/src/commands/call/contract.rs

@@ -5,6 +5,7 @@ use crate::{
 	common::{
 		builds::get_project_path,
 		contracts::has_contract_been_built,
+		prompt::display_message,
 		wallet::{prompt_to_use_wallet, request_signature},
 	},
 };
@@ -549,15 +550,6 @@ impl CallContractCommand {
 	}
 }
 
-fn display_message(message: &str, success: bool, cli: &mut impl Cli) -> Result<()> {
-	if success {
-		cli.outro(message)?;
-	} else {
-		cli.outro_cancel(message)?;
-	}
-	Ok(())
-}
-
 #[cfg(test)]
 mod tests {
 	use super::*;
@@ -1188,14 +1180,4 @@ mod tests {
 		assert!(!call_config.is_contract_build_required());
 		Ok(())
 	}
-
-	#[test]
-	fn display_message_works() -> Result<()> {
-		let mut cli = MockCli::new().expect_outro(&"Call completed successfully!");
-		display_message("Call completed successfully!", true, &mut cli)?;
-		cli.verify()?;
-		let mut cli = MockCli::new().expect_outro_cancel("Call failed.");
-		display_message("Call failed.", false, &mut cli)?;
-		cli.verify()
-	}
 }

crates/pop-cli/src/commands/mod.rs

@@ -5,6 +5,8 @@ use clap::Subcommand;
 use pop_common::templates::Template;
 use serde_json::{json, Value};
 
+#[cfg(feature = "parachain")]
+pub(crate) mod bench;
 pub(crate) mod build;
 pub(crate) mod call;
 pub(crate) mod clean;
@@ -23,6 +25,9 @@ pub(crate) enum Command {
 	#[clap(alias = "n")]
 	#[cfg(any(feature = "parachain", feature = "contract"))]
 	New(new::NewArgs),
+	/// Benchmark a pallet or parachain.
+	#[cfg(feature = "parachain")]
+	Bench(bench::BenchmarkArgs),
 	#[clap(alias = "b", about = about_build())]
 	#[cfg(any(feature = "parachain", feature = "contract"))]
 	Build(build::BuildArgs),
@@ -93,6 +98,8 @@ impl Command {
 					cmd.execute().await.map(|_| json!("default"))
 				},
 			},
+			#[cfg(feature = "parachain")]
+			Self::Bench(args) => bench::Command::execute(args).map(|_| Value::Null),
 			#[cfg(any(feature = "parachain", feature = "contract"))]
 			Self::Build(args) => match args.command {
 				None => build::Command::execute(args).map(|t| json!(t)),

crates/pop-cli/src/common/mod.rs

@@ -4,4 +4,5 @@ pub mod builds;
 #[cfg(feature = "contract")]
 pub mod contracts;
 pub mod helpers;
+pub mod prompt;
 pub mod wallet;

crates/pop-cli/src/common/prompt.rs

@@ -0,0 +1,29 @@
+use crate::cli::traits::Cli;
+use anyhow::Result;
+
+// Displays a message to the user, with formatting based on the success status.
+pub(crate) fn display_message(message: &str, success: bool, cli: &mut impl Cli) -> Result<()> {
+	if success {
+		cli.outro(message)?;
+	} else {
+		cli.outro_cancel(message)?;
+	}
+	Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+	use super::display_message;
+	use crate::cli::MockCli;
+	use anyhow::Result;
+
+	#[test]
+	fn display_message_works() -> Result<()> {
+		let mut cli = MockCli::new().expect_outro(&"Call completed successfully!");
+		display_message("Call completed successfully!", true, &mut cli)?;
+		cli.verify()?;
+		let mut cli = MockCli::new().expect_outro_cancel("Call failed.");
+		display_message("Call failed.", false, &mut cli)?;
+		cli.verify()
+	}
+}

crates/pop-cli/src/main.rs

@@ -1,60 +1,65 @@
 // SPDX-License-Identifier: GPL-3.0
 
 #[cfg(not(any(feature = "contract", feature = "parachain")))]
 compile_error!("feature \"contract\" or feature \"parachain\" must be enabled");
 
 use anyhow::{anyhow, Result};
 use clap::Parser;
 use commands::*;
 use serde_json::json;
 use std::{fs::create_dir_all, path::PathBuf};
 #[cfg(feature = "telemetry")]
 use {
 	pop_telemetry::{config_file_path, record_cli_command, record_cli_used, Telemetry},
 	std::env::args,
 };
 
 mod cli;
 #[cfg(any(feature = "parachain", feature = "contract"))]
 mod commands;
 mod common;
 mod style;
 mod wallet_integration;
 
 #[tokio::main]
 async fn main() -> Result<()> {
+	// Set environment for logging configuration, requires for `pop bench`.
+	if std::env::var("RUST_LOG").is_err() {
+		std::env::set_var("RUST_LOG", "info");
+	}
+
 	#[cfg(feature = "telemetry")]
 	let maybe_tel = init().unwrap_or(None);
 
 	let cli = Cli::parse();
 	let res = cli.command.execute().await;
 
 	#[cfg(feature = "telemetry")]
 	if let Some(tel) = maybe_tel.clone() {
 		// `args` is guaranteed to have at least 3 elements as clap will display help message if not
 		// set.
 		let (command, subcommand) = parse_args(args().collect());
 
 		if let Ok(sub_data) = &res {
 			// Best effort to send on first try, no action if failure.
 			let _ = record_cli_command(
 				tel.clone(),
 				&command,
 				json!({&subcommand: sub_data.to_string()}),
 			)
 			.await;
 		} else {
 			let _ = record_cli_command(tel, "error", json!({&command: &subcommand})).await;
 		}
 	}
 
 	// map result from Result<Value> to Result<()>
 	res.map(|_| ())
 }
 
 #[derive(Parser)]
 #[command(author, version, about, styles=style::get_styles())]
 pub struct Cli {
 	#[command(subcommand)]
 	command: Command,
 }

crates/pop-parachains/Cargo.toml

@@ -34,6 +34,12 @@ walkdir.workspace = true
 # Zombienet
 zombienet-sdk.workspace = true
 
+# Benchmarking
+cumulus-primitives-proof-size-hostfunction.workspace = true
+frame-benchmarking-cli.workspace = true
+sp-runtime.workspace = true
+sp-statement-store.workspace = true
+
 # Pop
 pop-common = { path = "../pop-common", version = "0.6.0" }
 

crates/pop-parachains/src/bench.rs

@@ -0,0 +1,17 @@
+use anyhow::Result;
+use frame_benchmarking_cli::PalletCmd;
+use sp_runtime::traits::BlakeTwo256;
+
+type HostFunctions = (
+	sp_statement_store::runtime_api::HostFunctions,
+	cumulus_primitives_proof_size_hostfunction::storage_proof_size::HostFunctions,
+);
+
+/// Generate benchmarks for a pallet.
+///
+/// # Arguments
+/// * `cmd` - Command to benchmark the extrinsic weight of FRAME Pallets.
+pub fn generate_benchmarks(cmd: &PalletCmd) -> Result<()> {
+	cmd.run_with_spec::<BlakeTwo256, HostFunctions>(None)
+		.map_err(|e| anyhow::anyhow!(format!("Failed to run benchmarking: {}", e.to_string())))
+}