feat: getNonce and setNonce cheatcodes

foundry-rs · Jan 30, 2022 · 3e24dc3 · 3e24dc3
1 parent 708f48f
commit 3e24dc3
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 0 deletions.
diff --git a/evm-adapters/src/sputnik/cheatcodes/cheatcode_handler.rs b/evm-adapters/src/sputnik/cheatcodes/cheatcode_handler.rs
@@ -867,6 +867,15 @@ 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::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: wait for https://github.com/rust-blockchain/evm/pull/92
+            }
         };
 
         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
@@ -179,6 +179,8 @@ pub enum CheatOp {
     CLEARMOCKEDCALLS,
     EXPECTCALL,
     GETCODE,
+    SETNONCE,
+    GETNONCE,
 }
 
 impl From<CheatOp> for OpCode {
@@ -212,6 +214,8 @@ impl CheatOp {
             CheatOp::CLEARMOCKEDCALLS => "VM_CLEARMOCKEDCALLS",
             CheatOp::EXPECTCALL => "VM_EXPECTCALL",
             CheatOp::GETCODE => "VM_GETCODE",
+            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
@@ -72,6 +72,8 @@ ethers::contract::abigen!(
             clearMockedCalls()
             expectCall(address,bytes)
             getCode(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
@@ -59,6 +59,12 @@ interface Hevm {
     function expectCall(address,bytes calldata) external;
 
     function getCode(string calldata) external returns (bytes memory);
+
+    // Get account nonce
+    function getNonce(address) external returns (uint256);
+
+    // Set account nonce
+    function setNonce(address,uint256) external;
 }
 
 contract HasStorage {
@@ -600,6 +606,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 RecordAccess {
@@ -773,6 +798,7 @@ contract ExpectEmit {
 }
 
 contract MockMe {
+    uint256 _a;
     function numberA() public returns (uint256) {
         return 1;
     }
@@ -784,6 +810,10 @@ contract MockMe {
     function add(uint256 a, uint256 b) public returns (uint256) {
         return a + b;
     }
+
+    function setA() public {
+        _a = 1;
+    }
 }
 
 contract MockInner {