feat(forge): getNonce and setNonce cheatcodes

foundry-rs · Feb 8, 2022 · ada53ee · ada53ee
1 parent 28f6d91
commit ada53ee
Show file tree

Hide file tree

Showing 5 changed files with 52 additions and 4 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/evm-adapters/src/sputnik/cheatcodes/cheatcode_handler.rs b/evm-adapters/src/sputnik/cheatcodes/cheatcode_handler.rs
@@ -867,13 +867,25 @@ impl<'a, 'b, B: Backend, P: PrecompileSet> CheatcodeStackExecutor<'a, 'b, B, P>
                 self.add_debug(CheatOp::EXPECTCALL);
                 self.state_mut().expected_calls.entry(inner.0).or_default().push(inner.1.to_vec());
             }
+
             HEVMCalls::Label(inner) => {
                 self.add_debug(CheatOp::LABEL);
                 let address = inner.0;
                 let label = inner.1;
 
                 self.state_mut().labels.insert(address, label);
             }
+
+            HEVMCalls::GetNonce(inner) => {
+                self.add_debug(CheatOp::GETNONCE);
+                let nonce = self.state().basic(inner.0).nonce;
+                res = ethers::abi::encode(&[Token::Uint(nonce)]);
+            }
+            HEVMCalls::SetNonce(inner) => {
+                self.add_debug(CheatOp::SETNONCE);
+                // TODO
+
+            }
         };
 
         self.fill_trace(&trace, true, Some(res.clone()), pre_index);

diff --git a/evm-adapters/src/sputnik/cheatcodes/debugger.rs b/evm-adapters/src/sputnik/cheatcodes/debugger.rs
@@ -180,6 +180,8 @@ pub enum CheatOp {
     EXPECTCALL,
     GETCODE,
     LABEL,
+    SETNONCE,
+    GETNONCE,
 }
 
 impl From<CheatOp> for OpCode {
@@ -214,6 +216,8 @@ impl CheatOp {
             CheatOp::EXPECTCALL => "VM_EXPECTCALL",
             CheatOp::GETCODE => "VM_GETCODE",
             CheatOp::LABEL => "VM_LABEL",
+            CheatOp::GETNONCE => "VM_GETNONCE",
+            CheatOp::SETNONCE => "VM_SETNONCE",
         }
     }
 }

diff --git a/evm-adapters/src/sputnik/cheatcodes/mod.rs b/evm-adapters/src/sputnik/cheatcodes/mod.rs
@@ -73,6 +73,8 @@ ethers::contract::abigen!(
             expectCall(address,bytes)
             getCode(string)
             label(address,string)
+            getNonce(address)
+            setNonce(address,uint256)
     ]"#,
 );
 pub use hevm_mod::{HEVMCalls, HEVM_ABI};

diff --git a/evm-adapters/testdata/CheatCodes.sol b/evm-adapters/testdata/CheatCodes.sol
@@ -61,6 +61,12 @@ interface Hevm {
     function getCode(string calldata) external returns (bytes memory);
     // Labels an address in call traces
     function label(address, string calldata) external;
+
+    // Get account nonce
+    function getNonce(address) external returns (uint256);
+
+    // Set account nonce
+    function setNonce(address,uint256) external;
 }
 
 contract HasStorage {
@@ -614,6 +620,25 @@ contract CheatCodes is DSTest {
             extcodecopy(who, add(o_code, 0x20), 0, size)
         }
     }
+
+    function testSetAndGetNonce() public {
+        address alice = address(1337);
+        assertEq(hevm.getNonce(alice), 0);
+
+        MockMe target = new MockMe();
+
+        hevm.startPrank(alice);
+        target.setA();
+        target.setA();
+        assertEq(hevm.getNonce(alice), 2);
+
+        hevm.setNonce(alice, 20);
+        assertEq(hevm.getNonce(alice), 20);
+
+        hevm.setNonce(alice, 14);
+        assertEq(hevm.getNonce(alice), 14);
+        hevm.stopPrank();
+    }
 }
 
 contract Label {
@@ -793,6 +818,7 @@ contract ExpectEmit {
 }
 
 contract MockMe {
+    uint256 _a;
     function numberA() public returns (uint256) {
         return 1;
     }
@@ -804,6 +830,10 @@ contract MockMe {
     function add(uint256 a, uint256 b) public returns (uint256) {
         return a + b;
     }
+
+    function setA() public {
+        _a = 1;
+    }
 }
 
 contract MockInner {