From 9a0f0c23ac8e9a96b9f36b5eab486d0579c12ffe Mon Sep 17 00:00:00 2001
From: Krishang <93703995+kamuik16@users.noreply.github.com>
Date: Wed, 10 Apr 2024 17:33:51 +0530
Subject: [PATCH] feat(forge): prompt address and uint cheatcodes (#7600)

* feat: prompt address and uint cheatcode

* nits

* chore: change test to pass ci
---
 crates/cheatcodes/assets/cheatcodes.json | 40 ++++++++++++++++++++++++
 crates/cheatcodes/spec/src/vm.rs         |  8 +++++
 crates/cheatcodes/src/fs.rs              | 16 ++++++++++
 testdata/cheats/Vm.sol                   |  2 ++
 testdata/default/cheats/Prompt.t.sol     | 11 +++++++
 5 files changed, 77 insertions(+)

diff --git a/crates/cheatcodes/assets/cheatcodes.json b/crates/cheatcodes/assets/cheatcodes.json
index 7a02c6ad64fa..c8b5d3d2affe 100644
--- a/crates/cheatcodes/assets/cheatcodes.json
+++ b/crates/cheatcodes/assets/cheatcodes.json
@@ -6231,6 +6231,26 @@
       "status": "stable",
       "safety": "safe"
     },
+    {
+      "func": {
+        "id": "promptAddress",
+        "description": "Prompts the user for an address in the terminal.",
+        "declaration": "function promptAddress(string calldata promptText) external returns (address);",
+        "visibility": "external",
+        "mutability": "",
+        "signature": "promptAddress(string)",
+        "selector": "0x62ee05f4",
+        "selectorBytes": [
+          98,
+          238,
+          5,
+          244
+        ]
+      },
+      "group": "filesystem",
+      "status": "stable",
+      "safety": "safe"
+    },
     {
       "func": {
         "id": "promptSecret",
@@ -6251,6 +6271,26 @@
       "status": "stable",
       "safety": "safe"
     },
+    {
+      "func": {
+        "id": "promptUint",
+        "description": "Prompts the user for uint256 in the terminal.",
+        "declaration": "function promptUint(string calldata promptText) external returns (uint256);",
+        "visibility": "external",
+        "mutability": "",
+        "signature": "promptUint(string)",
+        "selector": "0x652fd489",
+        "selectorBytes": [
+          101,
+          47,
+          212,
+          137
+        ]
+      },
+      "group": "filesystem",
+      "status": "stable",
+      "safety": "safe"
+    },
     {
       "func": {
         "id": "readCallers",
diff --git a/crates/cheatcodes/spec/src/vm.rs b/crates/cheatcodes/spec/src/vm.rs
index ae93475b2bf3..03a2cd090b85 100644
--- a/crates/cheatcodes/spec/src/vm.rs
+++ b/crates/cheatcodes/spec/src/vm.rs
@@ -1482,6 +1482,14 @@ interface Vm {
     #[cheatcode(group = Filesystem)]
     function promptSecret(string calldata promptText) external returns (string memory input);
 
+    /// Prompts the user for an address in the terminal.
+    #[cheatcode(group = Filesystem)]
+    function promptAddress(string calldata promptText) external returns (address);
+
+    /// Prompts the user for uint256 in the terminal.
+    #[cheatcode(group = Filesystem)]
+    function promptUint(string calldata promptText) external returns (uint256);
+
     // ======== Environment Variables ========
 
     /// Sets environment variables.
diff --git a/crates/cheatcodes/src/fs.rs b/crates/cheatcodes/src/fs.rs
index 3ba07b5b6cd2..9b22c4d8d2ae 100644
--- a/crates/cheatcodes/src/fs.rs
+++ b/crates/cheatcodes/src/fs.rs
@@ -1,6 +1,8 @@
 //! Implementations of [`Filesystem`](crate::Group::Filesystem) cheatcodes.
 
+use super::string::parse;
 use crate::{Cheatcode, Cheatcodes, Result, Vm::*};
+use alloy_dyn_abi::DynSolType;
 use alloy_json_abi::ContractObject;
 use alloy_primitives::U256;
 use alloy_sol_types::SolValue;
@@ -426,6 +428,20 @@ impl Cheatcode for promptSecretCall {
     }
 }
 
+impl Cheatcode for promptAddressCall {
+    fn apply(&self, state: &mut Cheatcodes) -> Result {
+        let Self { promptText: text } = self;
+        parse(&prompt(state, text, prompt_input)?, &DynSolType::Address)
+    }
+}
+
+impl Cheatcode for promptUintCall {
+    fn apply(&self, state: &mut Cheatcodes) -> Result {
+        let Self { promptText: text } = self;
+        parse(&prompt(state, text, prompt_input)?, &DynSolType::Uint(256))
+    }
+}
+
 pub(super) fn write_file(state: &Cheatcodes, path: &Path, contents: &[u8]) -> Result {
     let path = state.config.ensure_path_allowed(path, FsAccessKind::Write)?;
     // write access to foundry.toml is not allowed
diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol
index 1d0b6118b1a3..44f0d52ac423 100644
--- a/testdata/cheats/Vm.sol
+++ b/testdata/cheats/Vm.sol
@@ -307,7 +307,9 @@ interface Vm {
     function prevrandao(bytes32 newPrevrandao) external;
     function projectRoot() external view returns (string memory path);
     function prompt(string calldata promptText) external returns (string memory input);
+    function promptAddress(string calldata promptText) external returns (address);
     function promptSecret(string calldata promptText) external returns (string memory input);
+    function promptUint(string calldata promptText) external returns (uint256);
     function readCallers() external returns (CallerMode callerMode, address msgSender, address txOrigin);
     function readDir(string calldata path) external view returns (DirEntry[] memory entries);
     function readDir(string calldata path, uint64 maxDepth) external view returns (DirEntry[] memory entries);
diff --git a/testdata/default/cheats/Prompt.t.sol b/testdata/default/cheats/Prompt.t.sol
index dadfd30a97b2..9e461c2b527e 100644
--- a/testdata/default/cheats/Prompt.t.sol
+++ b/testdata/default/cheats/Prompt.t.sol
@@ -3,6 +3,7 @@ pragma solidity 0.8.18;
 
 import "ds-test/test.sol";
 import "cheats/Vm.sol";
+import "../logs/console.sol";
 
 contract PromptTest is DSTest {
     Vm constant vm = Vm(HEVM_ADDRESS);
@@ -15,4 +16,14 @@ contract PromptTest is DSTest {
         vm._expectCheatcodeRevert();
         vm.promptSecret("test");
     }
+
+    function testPrompt_Address() public {
+        vm._expectCheatcodeRevert();
+        address test = vm.promptAddress("test");
+    }
+
+    function testPrompt_Uint() public {
+        vm._expectCheatcodeRevert();
+        uint256 test = vm.promptUint("test");
+    }
 }