From 306a6d3c23e381fb238ff0114b7a2206272f2a73 Mon Sep 17 00:00:00 2001
From: cygaar <tmpftl1@gmail.com>
Date: Sun, 14 Apr 2024 22:49:52 -0400
Subject: [PATCH] Add features flag for cargo stylus check/deploy

---
 src/check.rs   | 1 +
 src/deploy.rs  | 1 +
 src/main.rs    | 3 +++
 src/project.rs | 7 +++++++
 4 files changed, 12 insertions(+)

diff --git a/src/check.rs b/src/check.rs
index 9270e06..202f613 100644
--- a/src/check.rs
+++ b/src/check.rs
@@ -65,6 +65,7 @@ pub async fn run_checks(cfg: CheckConfig) -> eyre::Result<bool> {
         Some(path) => PathBuf::from_str(path).unwrap(),
         None => project::build_project_dylib(BuildConfig {
             opt_level: project::OptLevel::default(),
+            features: cfg.features.clone(),
             nightly: cfg.nightly,
             rebuild: true,
             skip_contract_size_check: cfg.skip_contract_size_check,
diff --git a/src/deploy.rs b/src/deploy.rs
index 9f68b63..4a3492f 100644
--- a/src/deploy.rs
+++ b/src/deploy.rs
@@ -116,6 +116,7 @@ programs to Stylus chains here https://docs.arbitrum.io/stylus/stylus-quickstart
             Some(path) => PathBuf::from_str(path).unwrap(),
             None => project::build_project_dylib(BuildConfig {
                 opt_level: project::OptLevel::default(),
+                features: cfg.check_cfg.features,
                 nightly: cfg.check_cfg.nightly,
                 rebuild: false, // The check step at the start of this command rebuilt.
                 skip_contract_size_check: cfg.check_cfg.skip_contract_size_check,
diff --git a/src/main.rs b/src/main.rs
index d92ea9c..44304ce 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -95,6 +95,9 @@ pub struct CheckConfig {
     /// Specifies a WASM file instead of looking for one in the current directory.
     #[arg(long)]
     wasm_file_path: Option<String>,
+    /// Specifies the features to use when building the Stylus binary.
+    #[arg(long)]
+    features: Option<String>,
     /// Specify the program address we want to check activation for. If unspecified, it will
     /// compute the next program address from the user's wallet address and nonce, which will require
     /// wallet-related flags to be specified.
diff --git a/src/project.rs b/src/project.rs
index 2188bf4..f22d544 100644
--- a/src/project.rs
+++ b/src/project.rs
@@ -21,6 +21,7 @@ pub enum OptLevel {
 
 pub struct BuildConfig {
     pub opt_level: OptLevel,
+    pub features: Option<String>,
     pub nightly: bool,
     pub rebuild: bool,
     pub skip_contract_size_check: bool,
@@ -65,6 +66,10 @@ pub fn build_project_dylib(cfg: BuildConfig) -> Result<PathBuf> {
         cmd.arg("build");
         cmd.arg("--lib");
 
+        if cfg.features != None {
+            cmd.arg(format!("--features={}", cfg.features.clone().unwrap()));
+        }
+
         if cfg.nightly {
             cmd.arg("-Z");
             cmd.arg("build-std=std,panic_abort");
@@ -126,6 +131,7 @@ https://github.com/OffchainLabs/cargo-stylus/blob/main/OPTIMIZING_BINARIES.md"#,
                     // Attempt to build again with a bumped-up optimization level.
                     return build_project_dylib(BuildConfig {
                         opt_level: OptLevel::Z,
+                        features: cfg.features,
                         nightly: cfg.nightly,
                         rebuild: true,
                         skip_contract_size_check: cfg.skip_contract_size_check,
@@ -141,6 +147,7 @@ https://github.com/OffchainLabs/cargo-stylus/blob/main/OPTIMIZING_BINARIES.md"#,
                         // only available with nightly compilation.
                         return build_project_dylib(BuildConfig {
                             opt_level: OptLevel::Z,
+                            features: cfg.features,
                             nightly: true,
                             rebuild: true,
                             skip_contract_size_check: cfg.skip_contract_size_check,