feat(cheatcode): impl getDeployedCode to handle immutable variables

[xenoliss]

Motivation

While writing forge unit tests it's often that I need to mock some contracts at specific addresses. For this I'm using 2 cheatcodes as shown below:

bytes memory patchedCode = vm.getDeployedCode("ERC20.sol:ERC20");
vm.etch(target, patchedCode);

This works well when the contract bytecode does not include immutable variables.

Using the code above with a contract that defines immutable variables will let the immutable references uninitialized in the contract deployed bytecode. Most of the time this makes the contract unusable if the immutable variables are used by the methods we plan to call.

Solution

This PR extends the current getDeployedCode cheatcode by adding another prototype which accepts a list of bytes32 values. Those bytes32 values are given to the get_deployed_code method and passed to the into_deployed_bytecode trait method. Inside the into_deployed_bytecode method, the ArtifactBytecode::Forge handles the case where there are some immutable variables defined and if so, ensures their values are specified and manually patches the byte code to set their values.

NOTE 1: The implementation assumes that the bytes length of an immutable variable is 32.
NOTE 2: I had to move ArtifactBytecode::Forge first in the ArtifactBytecode enum to give it the priority when deserializing.
NOTE 3: Manually patching (with a nested loop) the bytecode is inefficient and I'm open to suggestions to improve this.
NOTE 4: into_deployed_bytecode was returning an Option so I kept it this way but it may be preferable now to have it return a Result<Option<Bytes>, Error> with Error used to describe an problem related to immutable variables initialization.

Here is an example of the new getDeployedCode prototype usage:

// SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.21;

import {Test, console} from "forge-std/Test.sol";

contract Immutable {
    uint256 public immutable a;
    uint256 public immutable b;
    bool public immutable c;

    constructor(uint256 _a, uint256 _b, bool _c) {
        a = _a;
        b = _b;
        c = _c;
    }
}

contract TestGetDeployedCode is Test {
    // Test fails because we're trying to deploy a ccontract with immutabel variables without
    // specifying their individual values.
    function testFail_previous() public {
        address previous = makeAddr("previous");

        bytes memory patchedCode = vm.getDeployedCode("TestMe.t.sol:Immutable");
    }

    function test_next() public {
        address next = makeAddr("next");

        uint256 a = 1;
        uint256 b = 2;
        bool c = true;
        bytes32[] memory immutables = new bytes32[](3);
        immutables[0] = bytes32(uint256(a));
        immutables[1] = bytes32(uint256(b));
        immutables[2] = bytes32(uint256(c ? 1 : 0));

        bytes memory patchedCode = vm.getDeployedCode(
            "TestMe.t.sol:Immutable",
            immutables
        );
        vm.etch(next, patchedCode);

        console.log("a", Immutable(next).a());
        console.log("b", Immutable(next).b());
        console.log("c", Immutable(next).c());
    }
}

Which outputs:

$ forge test -vv 
[⠑] Compiling...
No files changed, compilation skipped

Running 2 tests for test/TestMe.t.sol:TestGetDeployedCode
[PASS] testFail_previous() (gas: 4590)
[PASS] test_next() (gas: 14108)
Logs:
  a 1
  b 2
  c true

Test result: ok. 2 passed; 0 failed; 0 skipped; finished in 2.16ms
Ran 1 test suites: 2 tests passed, 0 failed, 0 skipped (2 total tests)

[xenoliss]

Noticed there is a problem now with the test_artifact_parsing_no_immutables, certainly due to the ArtifactBytecode enum reordering. Will work on that when I have time.

[mds1]

This is neat, can you expand more on the UX here? For example, given this contract:

contract Counter {
    uint256 public number;
    uint256 public immutable immutable1;
    uint256 public immutable immutable2;

    constructor(uint256 x) {
        immutable1 = x;
        immutable2 = x * 2;
    }

    function foo(uint256 newNumber) public {
        number = newNumber + immutable1;
    }

    function bar() public {
        number += immutable2;
    }
}

The compiler outputs:

"immutableReferences": {
  "5": [
    {
      "start": 184,
      "length": 32
    },
    {
      "start": 231,
      "length": 32
    }
  ],
  "7": [
    {
      "start": 97,
      "length": 32
    },
    {
      "start": 276,
      "length": 32
    }
  ]
}

As a user, the UX is:

• For each immutable reference, take the the key and use that as the ID in the AST to find the immutable's name.
• Then the array passed to getDeployedCode has length 2 since immutableReferences
• And the entries should be in order of increasing ID

Is that correct?

[xenoliss]

Your comment made me realize I did not handle the case where an immutable variable was referenced several times in the contract byte code. Commit 35060f14 should fix this (and clean a bit the code).

Regarding the UX:

For each immutable reference, take the the key and use that as the ID in the AST to find the immutable's name.

This is I guess the best way to make sure we're setting the right value for the right immutable variable but as far as I've seen, immutable var declaration is in order so the user might also use the declaration order to set the values (special care must be taken in case of inheritance).

Then the array passed to getDeployedCode has length 2 since immutableReferences

Yes in the given example as there are 2 immutable variable in the contract the array passed to the getDeployedCode should only contains 2 values. If this is not the case None will be returned here.

And the entries should be in order of increasing ID

Right, if I have immutable1, immutable2, immutable3 with AST id 3, 7 and 5, the user must fill the immutables array with [immutable1Value, immutable3Value, immutable2Value].

NOTE: As mentioned here, I expect that iterating over the immutable_references BTreeMap will yield the keys (corresponding to the AST ids of the immutable variables) in ascending order. This is mandatory for the current implementation to work and as far as I've seen this is always the case. If you know a case where this will not hold a fix will be needed.

REMANING ISSUE: I'm not 100% sure to understand how the ArtifactBytecode deserialization works.
Moving the Forge makes the test_artifact_parsing_with_immutables test to pass but test_artifact_parsing_no_immutables fails. Contrary, keeping Forge variant in 3rd position makes test_artifact_parsing_no_immutables but test_artifact_parsing_with_immutables fails. I'm not sure how to properly guide serde into deserializing into the best variant.

[mattsse]

some nits

still need to catch up with the discussion

[mattsse]

I don't know what this does.

this needs docs, I'd be totally lost when I have to revisit this again

[mattsse]

I don't really want to modify these functions,

but I'm open to additional functions _with_immutables or something like this

[mattsse]

Moving the Forge makes the test_artifact_parsing_with_immutables test to pass but

what does this mean?

[xenoliss]

Moving the Forge makes the test_artifact_parsing_with_immutables test to pass but

Not well explained indeed.

I meant that I'm struggling to have serde figure out what variant of the ArtifactBytecode enum the json payload must be deserialized into. I tried to re-order the variants by moving the Forge variant up in the enum but this makes the test_artifact_parsing_no_immutables fail because it also tries to deserialize the payload into a Forge variant.

[mattsse]

could you add the json artifacts, how do they differ?

[xenoliss]

The json artifacts are those I added in the test-data. There is a solc and forge versions that I'm using in the unit tests.

[Evalir]

just catching up—wondering:

• Is expecting the variables to be 32 bytes long a reasonable invariant to hold? When does this break down? I'm feeling this could be a huge footgun.
• What would be the recommended way to deal with inheritance? I feel this could get messy

[mds1]

Is expecting the variables to be 32 bytes long a reasonable invariant to hold? When does this break down? I'm feeling this could be a huge footgun.

Currently solidity does not allow immutables for dynamic length types, and all immutables (even smaller types like uint8) are padded to the full 32 bytes, so this is ok for now. But solidity will likely lift this restrictions—if you click the link in this comment it seems it's planned for Q1 2024

What would be the recommended way to deal with inheritance? I feel this could get messy

Agreed, I think it's the workflow I explained here but @xenoliss can confirm. If so, let's make sure that workflow gets documented in the book: #5718 (comment)

[Evalir]

Currently solidity does not allow immutables for dynamic length types, and all immutables (even smaller types like uint8) are padded to the full 32 bytes, so this is ok for now. But solidity will likely lift this restrictions—if you click the link in ethereum/solidity#12587 (comment) it seems it's planned for Q1 2024

Yeah this was my worry—this approach seems good for now but not super future proof. We can probably go with this for now and iterate later, but we def should keep track of this so it doesn't bite people later.

Agreed, I think it's the workflow I explained here but @xenoliss can confirm. If so, let's make sure that workflow gets documented in the book: #5718 (comment)

gotcha, cool!

[xenoliss]

Agreed, I think it's the workflow I explained here but @xenoliss can confirm.

Yes it is, I also explained it here.

Regarding inheritance the user must indeed be aware of all the immutable variables brought by inherited contracts. In case the user is not providing all the immutable values the current implementation will return None and return a failling message (could be improved).

Regarding the way immutable values are currently padded to 32 bytes it's indeed not future proofed. If you have other ideas I have no problem to change the interface (also I feel like it's better to have it, knowing this restriction, rather than not having it at all).

The main technical problem remaining on my side is the deserialization of the json payload into the ArtifactBytecode enum which sometimes pick the wrong variant (see the unit tests).

[DaniPopes]

Hey @xenoliss, we recently refactored how cheatcodes are implemented in #6131. You can read more about it in the dev docs. Please let me know if you have any problems rebasing!

[zerosnacks]

Hi @xenoliss, thanks for your PR! Would be great to get this merged - would you be interested in picking this back up? The new cheatcode system is really neat. If not I can also take it from here.