feat(cheatcodes): vm.mockCalls() for mocking multiple calls

[Tudmotu]

Component

Forge

Describe the feature you would like

function mockCalls(address where, bytes calldata data, bytes[] calldata retdata) external;
function mockCalls(
    address where,
    uint256 value,
    bytes calldata data,
    bytes[] calldata retdata
) external;

Notice the bytes[] for the retdata. The idea is to mock the same call with different return values each time.

Additional context

Motivation

Today, we can mock the return value of a function call by using mockCall. Sometimes, it is useful to mock the same call twice with different return values.

Consider the following function:

function failsafeArb (IERC20 token, IERC20 arbToken) public {
    uint initialBalance = token.balanceOf(address(this));
    doArbitrage(token, arbToken);
    uint finalBalance = token.balanceOf(address(this));
    require(finalBalance > initialBalance, 'arbitrage failed');
}

Both invocations of token.balanceOf() have the same calldata. This means we cannot use two different mockCall to test this behavior.

With mockCalls we would test it like so:

bytes[] retdata = new bytes[](2);
retdata[0] = abi.encode(100 ether);
retdata[1] = abi.encode(99 ether);
vm.mockCalls(
    address(token),
    abi.encodeWithSelector(
        IERC20.balanceOf.selector,
        address(testContract)
    ),
    retdata
);

vm.expectRevert('arbitrage failed');
testContract.failsafeArb(token, arbToken);

Suggested implementation

1. Change the mock_calls field in Cheatcodes to:

pub mocked_calls: HashMap<Address, BTreeMap<MockCallDataContext, Vec<MockCallReturnData>>>

2. To avoid touching existing cheatcode impls, modify mock_call in mock.rs to vec![] the return data:

state.mocked_calls.entry(*callee).or_default().insert(
    MockCallDataContext { calldata: Bytes::copy_from_slice(cdata), value: value.copied() },
    vec![MockCallReturnData { ret_type, data: Bytes::copy_from_slice(rdata) }],
);

3. In inspector.rs in the call function, use the MockCallReturnData vector as a stack, removing the first item each time the inspector is triggered
4. Add the mockCalls cheatcodes, including a mock_calls function that takes rdata: Vec<Vec<u8>> as argument and maps it to a MockCallReturnData vector

Open Qs

1. What should happen if we mock 2 calls but invoke the function 3 times? Should the last retdata stay for the rest of the transaction (i.e. the stack not be emptied, but kept at size 1)? This would preserve existing behavior of mockCall
2. The Solidity API in the example I provided above with a dynamic array feels a bit clunky. Does anyone have an idea for a better API?

Notes

I did not try to implement this yet. This is just from looking at the code. First I'd like to know if this is something you'd accept a PR for. I will of course take it on myself. If this feature is desired, we can better define the API and the implementation.

Thank you 🙂

[grandizzy]

@Tudmotu thank you for the report! this improvement is not on our roadmap for 1.0 but we would be happy to review and accept PR for (and include it). thank you!

[Tudmotu]

@grandizzy I could take a stab at it though I'm not sure if the proposed design is still relevant since I proposed this a while ago. I will take a look.

[Tudmotu]

@grandizzy seems like the design is still applicable, I opened a PR: #9024