feat: Refactor assumeNot* std cheats

src/StdCheats.sol

@@ -9,6 +9,9 @@ import {Vm} from "./Vm.sol";
 abstract contract StdCheatsSafe {
     Vm private constant vm = Vm(address(uint160(uint256(keccak256("hevm cheat code")))));
 
+    uint256 private constant UINT256_MAX =
+        115792089237316195423570985008687907853269984665640564039457584007913129639935;
+
     bool private gasMeteringOff;
 
     // Data structures to parse Transaction objects from the broadcast artifact
@@ -193,6 +196,14 @@ abstract contract StdCheatsSafe {
         uint256 key;
     }
 
+    enum AddressType {
+        Payable,
+        NonPayable,
+        ZeroAddress,
+        Precompiles,
+        ForgeAddresses
+    }
+
     // Checks that `addr` is not blacklisted by token contracts that have a blacklist.
     function assumeNotBlacklisted(address token, address addr) internal view virtual {
         // Nothing to check if `token` is not a contract.
@@ -222,6 +233,84 @@ abstract contract StdCheatsSafe {
         assumeNotBlacklisted(token, addr);
     }
 
+    function assumeAddressIsNot(AddressType addressType, address addr) internal virtual {
+        if (addressType == AddressType.Payable) {
+            assumeNoPayable(addr);
+        } else if (addressType == AddressType.NonPayable) {
+            assumeNoNonPayable(addr);
+        } else if (addressType == AddressType.ZeroAddress) {
+            assumeNoZeroAddress(addr);
+        } else if (addressType == AddressType.Precompiles) {
+            assumeNoPrecompiles(addr);
+        } else if (addressType == AddressType.ForgeAddresses) {
+            assumeNoForgeAddresses(addr);
+        }
+    }
+
+    function assumeAddressIsNot(AddressType addressType1, AddressType addressType2, address addr) internal virtual {
+        assumeAddressIsNot(addressType1, addr);
+        assumeAddressIsNot(addressType2, addr);
+    }
+
+    function assumeAddressIsNot(
+        AddressType addressType1,
+        AddressType addressType2,
+        AddressType addressType3,
+        address addr
+    ) internal virtual {
+        assumeAddressIsNot(addressType1, addr);
+        assumeAddressIsNot(addressType2, addr);
+        assumeAddressIsNot(addressType3, addr);
+    }
+
+    function assumeAddressIsNot(
+        AddressType addressType1,
+        AddressType addressType2,
+        AddressType addressType3,
+        AddressType addressType4,
+        address addr
+    ) internal virtual {
+        assumeAddressIsNot(addressType1, addr);
+        assumeAddressIsNot(addressType2, addr);
+        assumeAddressIsNot(addressType3, addr);
+        assumeAddressIsNot(addressType4, addr);
+    }
+
+    function _isPayable(address addr) private returns (bool) {
+        uint256 size;
+        assembly {
+            size := extcodesize(addr)
+        }
+
+        if (size == 0) {
+            // return false if no code
+            return false;
+        } else {
+            require(addr.balance < UINT256_MAX, "balance exceeds max uint256");
+            uint256 origBalanceTest = address(this).balance;
+            uint256 origBalanceAddr = address(addr).balance;
+            (bool success,) = payable(addr).call{value: 1}("");
+
+            // reset balances
+            vm.deal(address(this), origBalanceTest);
+            vm.deal(addr, origBalanceAddr);
+
+            return success;
+        }
+    }
+
+    function assumeNoPayable(address addr) internal virtual {
+        vm.assume(!_isPayable(addr));
+    }
+
+    function assumeNoNonPayable(address addr) internal virtual {
+        vm.assume(_isPayable(addr));
+    }
+
+    function assumeNoZeroAddress(address addr) internal pure virtual {
+        vm.assume(addr != address(0));
+    }
+
     function assumeNoPrecompiles(address addr) internal pure virtual {
         assumeNoPrecompiles(addr, _pureChainId());
     }
@@ -249,6 +338,13 @@ abstract contract StdCheatsSafe {
         // forgefmt: disable-end
     }
 
+    function assumeNoForgeAddresses(address addr) internal pure virtual {
+        // vm and console addresses
+        vm.assume(
+            addr != 0x7109709ECfa91a80626fF3989D68f67F5b1DD12D || addr != 0x000000000000000000636F6e736F6c652e6c6f67
+        );
+    }
+
     function readEIP1559ScriptArtifact(string memory path)
         internal
         view
@@ -512,13 +608,6 @@ abstract contract StdCheatsSafe {
         }
     }
 
-    // a cheat for fuzzing addresses that are payable only
-    // see https://github.com/foundry-rs/foundry/issues/3631
-    function assumePayable(address addr) internal virtual {
-        (bool success,) = payable(addr).call{value: 0}("");
-        vm.assume(success);
-    }
-
     // We use this complex approach of `_viewChainId` and `_pureChainId` to ensure there are no
     // compiler warnings when accessing chain ID in any solidity version supported by forge-std. We
     // can't simply access the chain ID in a normal view or pure function because the solc View Pure

test/StdCheats.t.sol

@@ -335,41 +335,52 @@ contract StdCheatsTest is Test {
         return number;
     }
 
-    function testAssumeNoPrecompiles(address addr) external {
-        assumeNoPrecompiles(addr, getChain("optimism_goerli").chainId);
+    function testAssumeAddressIsNot(address addr) external {
+        // skip over Payable and NonPayable enums
+        for (uint8 i = 2; i < uint8(type(AddressType).max); i++) {
+            assumeAddressIsNot(AddressType(i), addr);
+        }
+        assertTrue(addr != address(0));
+        assertTrue(addr < address(1) || addr > address(9));
         assertTrue(
-            addr < address(1) || (addr > address(9) && addr < address(0x4200000000000000000000000000000000000000))
-                || addr > address(0x4200000000000000000000000000000000000800)
+            addr != 0x7109709ECfa91a80626fF3989D68f67F5b1DD12D || addr != 0x000000000000000000636F6e736F6c652e6c6f67
         );
     }
 
-    function _assumePayable(address addr) public {
-        assumePayable(addr);
+    function testAssumeNoPayable() external {
+        // all should pass since these addresses are not payable
+
+        // VM address
+        assumeNoPayable(0x7109709ECfa91a80626fF3989D68f67F5b1DD12D);
+
+        // Console address
+        assumeNoPayable(0x000000000000000000636F6e736F6c652e6c6f67);
+
+        // Create2Deployer
+        assumeNoPayable(0x4e59b44847b379578588920cA78FbF26c0B4956C);
     }
 
-    function testAssumePayable() external {
-        // We deploy a mock version so we can properly test the revert.
-        StdCheatsMock stdCheatsMock = new StdCheatsMock();
+    function testAssumeNoNonPayable() external {
         // all should revert since these addresses are not payable
 
         // VM address
         vm.expectRevert();
-        stdCheatsMock.exposed_assumePayable(0x7109709ECfa91a80626fF3989D68f67F5b1DD12D);
+        assumeNoNonPayable(0x7109709ECfa91a80626fF3989D68f67F5b1DD12D);
 
         // Console address
         vm.expectRevert();
-        stdCheatsMock.exposed_assumePayable(0x000000000000000000636F6e736F6c652e6c6f67);
+        assumeNoNonPayable(0x000000000000000000636F6e736F6c652e6c6f67);
 
         // Create2Deployer
         vm.expectRevert();
-        stdCheatsMock.exposed_assumePayable(0x4e59b44847b379578588920cA78FbF26c0B4956C);
+        assumeNoNonPayable(0x4e59b44847b379578588920cA78FbF26c0B4956C);
     }
 
-    function testAssumePayable(address addr) external {
-        assumePayable(addr);
+    function testAssumeNoPrecompiles(address addr) external {
+        assumeNoPrecompiles(addr, getChain("optimism_goerli").chainId);
         assertTrue(
-            addr != 0x7109709ECfa91a80626fF3989D68f67F5b1DD12D && addr != 0x000000000000000000636F6e736F6c652e6c6f67
-                && addr != 0x4e59b44847b379578588920cA78FbF26c0B4956C
+            addr < address(1) || (addr > address(9) && addr < address(0x4200000000000000000000000000000000000000))
+                || addr > address(0x4200000000000000000000000000000000000800)
         );
     }
 
@@ -402,10 +413,6 @@ contract StdCheatsTest is Test {
 }
 
 contract StdCheatsMock is StdCheats {
-    function exposed_assumePayable(address addr) external {
-        assumePayable(addr);
-    }
-
     // We deploy a mock version so we can properly test expected reverts.
     function exposed_assumeNotBlacklisted(address token, address addr) external view {
         return assumeNotBlacklisted(token, addr);