From 373ad46de9034f3b9e30b95084c9d1bd076d66a7 Mon Sep 17 00:00:00 2001
From: Yash Atreya <44857776+yash-atreya@users.noreply.github.com>
Date: Wed, 9 Oct 2024 11:31:38 +0530
Subject: [PATCH] feat(`cheatcodes`): vm.getScriptWallets() (#9052)

* feat(`cheatcodes`): vm.getScriptWallets()

* feat: load default anvil accounts in script

* Revert "feat: load default anvil accounts in script"

This reverts commit 4d64356a51bf226482269a2af47f947c4e49e462.

* clippy

* test

---------

Co-authored-by: grandizzy <grandizzy.the.egg@gmail.com>

---------

Co-authored-by: grandizzy <grandizzy.the.egg@gmail.com>
---
 crates/cheatcodes/assets/cheatcodes.json | 20 ++++++++++++
 crates/cheatcodes/spec/src/vm.rs         |  4 +++
 crates/cheatcodes/src/script.rs          | 15 +++++++++
 crates/forge/tests/cli/script.rs         | 40 ++++++++++++++++++++++++
 testdata/cheats/Vm.sol                   |  1 +
 5 files changed, 80 insertions(+)

diff --git a/crates/cheatcodes/assets/cheatcodes.json b/crates/cheatcodes/assets/cheatcodes.json
index e6fef48b0891..06ddba5bd5d1 100644
--- a/crates/cheatcodes/assets/cheatcodes.json
+++ b/crates/cheatcodes/assets/cheatcodes.json
@@ -5395,6 +5395,26 @@
       "status": "stable",
       "safety": "safe"
     },
+    {
+      "func": {
+        "id": "getScriptWallets",
+        "description": "Returns addresses of available unlocked wallets in the script environment.",
+        "declaration": "function getScriptWallets() external returns (address[] memory wallets);",
+        "visibility": "external",
+        "mutability": "",
+        "signature": "getScriptWallets()",
+        "selector": "0x7c49aa1f",
+        "selectorBytes": [
+          124,
+          73,
+          170,
+          31
+        ]
+      },
+      "group": "scripting",
+      "status": "stable",
+      "safety": "safe"
+    },
     {
       "func": {
         "id": "indexOf",
diff --git a/crates/cheatcodes/spec/src/vm.rs b/crates/cheatcodes/spec/src/vm.rs
index 0ee95e43fd02..e73755de1d3b 100644
--- a/crates/cheatcodes/spec/src/vm.rs
+++ b/crates/cheatcodes/spec/src/vm.rs
@@ -1877,6 +1877,10 @@ interface Vm {
     #[cheatcode(group = Scripting)]
     function broadcastRawTransaction(bytes calldata data) external;
 
+    /// Returns addresses of available unlocked wallets in the script environment.
+    #[cheatcode(group = Scripting)]
+    function getScriptWallets() external returns (address[] memory wallets);
+
     // ======== Utilities ========
 
     // -------- Strings --------
diff --git a/crates/cheatcodes/src/script.rs b/crates/cheatcodes/src/script.rs
index 93d5aaaf8d34..f9535844c7c4 100644
--- a/crates/cheatcodes/src/script.rs
+++ b/crates/cheatcodes/src/script.rs
@@ -3,6 +3,7 @@
 use crate::{Cheatcode, CheatsCtxt, Result, Vm::*};
 use alloy_primitives::{Address, B256, U256};
 use alloy_signer_local::PrivateKeySigner;
+use alloy_sol_types::SolValue;
 use foundry_wallets::{multi_wallet::MultiWallet, WalletSigner};
 use parking_lot::Mutex;
 use std::sync::Arc;
@@ -60,6 +61,20 @@ impl Cheatcode for stopBroadcastCall {
     }
 }
 
+impl Cheatcode for getScriptWalletsCall {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let script_wallets =
+            ccx.state.script_wallets().cloned().map(|sw| sw.signers().unwrap_or_default());
+
+        if let Some(script_wallets) = script_wallets {
+            let script_wallets: Vec<Address> = script_wallets.into_iter().collect();
+            Ok(script_wallets.abi_encode())
+        } else {
+            Ok(Default::default())
+        }
+    }
+}
+
 #[derive(Clone, Debug, Default)]
 pub struct Broadcast {
     /// Address of the transaction origin
diff --git a/crates/forge/tests/cli/script.rs b/crates/forge/tests/cli/script.rs
index 61c2456fef66..826d66282dd7 100644
--- a/crates/forge/tests/cli/script.rs
+++ b/crates/forge/tests/cli/script.rs
@@ -2055,3 +2055,43 @@ ONCHAIN EXECUTION COMPLETE & SUCCESSFUL.
 
 "#]]);
 });
+
+forgetest_init!(can_get_script_wallets, |prj, cmd| {
+    let script = prj
+        .add_source(
+            "Foo",
+            r#"
+import "forge-std/Script.sol";
+
+interface Vm {
+    function getScriptWallets() external returns (address[] memory wallets);
+}
+
+contract WalletScript is Script {
+    function run() public {
+        address[] memory wallets = Vm(address(vm)).getScriptWallets();
+        console.log(wallets[0]);
+    }
+}"#,
+        )
+        .unwrap();
+    cmd.arg("script")
+        .arg(script)
+        .args([
+            "--private-key",
+            "0x2a871d0798f97d79848a013d4936a73bf4cc922c825d33c1cf7073dff6d409c6",
+            "-v",
+        ])
+        .assert_success()
+        .stdout_eq(str![[r#"
+[COMPILING_FILES] with [SOLC_VERSION]
+[SOLC_VERSION] [ELAPSED]
+Compiler run successful!
+Script ran successfully.
+[GAS]
+
+== Logs ==
+  0xa0Ee7A142d267C1f36714E4a8F75612F20a79720
+
+"#]]);
+});
diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol
index be6522ac44ed..1458e3e4621d 100644
--- a/testdata/cheats/Vm.sol
+++ b/testdata/cheats/Vm.sol
@@ -265,6 +265,7 @@ interface Vm {
     function getNonce(address account) external view returns (uint64 nonce);
     function getNonce(Wallet calldata wallet) external returns (uint64 nonce);
     function getRecordedLogs() external returns (Log[] memory logs);
+    function getScriptWallets() external returns (address[] memory wallets);
     function indexOf(string calldata input, string calldata key) external pure returns (uint256);
     function isContext(ForgeContext context) external view returns (bool result);
     function isDir(string calldata path) external returns (bool result);