Add mapping slot API to cheatcodes

[clouds56]

Motivation

When testing a contract, we have no knowledge about how many elements there're in a mapping, so we'd like to add a family of cheatcode to solve this problem.

Solution

1. we add these api to Cheatcodes

            startMappingRecording()
            getMappingLength(address,bytes32)
            getMappingSlotAt(address,bytes32,uint256)
            getMappingKeyOf(address,bytes32)
            getMappingParentOf(address,bytes32)

2. after startMappingRecording() we record all potential map sstore to Cheatcodes local state
3. getMappingLength(address target, bytes32 mappingSlot) returns (uint256) would get count of elements in mappingSlot at target address.
4. getMappingSlotAt(address target, bytes32 mappingSlot, uint256 idx) returns (bytes) would get elements of idx in the mapping 0 <= idx < length, the length is returned by getMappingLength
5. getMappingKeyOf(address target, bytes32 elementSlot) returns (bytes) and getMappingParentOf(address target, bytes32 elementSlot) returns (bytes) are called on "element slot" returned by getMappingSlotAt

Details

• we would record on sha3 and sstore on step
• when see a slot in sstore, we would first check if it is generated by some sha3. If so, this would probable be a element of some mapping.

[mattsse]

When testing a contract, we have no knowledge about how many elements there're in a mapping,
Seems like this would be a useful cheatcode

wdyt @mds1 ?

[mattsse]

test and code look great.

only smol nits

[flyq]

issue:
#4635

more test:
https://github.com/0xevm/sloads_demo/

[mattsse]

lgtm, pending @mds1

[mds1]

This is a really interesting idea, thanks @clouds56! Being able to loop over mapping keys is very useful, especially during invariant tests, so people often use arrays to manually track all keys written to a mapping to work around this limitation.

I want to think on this API a bit before merging, as I'm not yet convinced it's the best UX. Some scattered initial thoughts for now:

• Having to use slot numbers isn't ideal since they can change as contract development progresses
• Having to use slot numbers isn't ideal since it's not very user friendly, it would be nice if you can reference mappings by their variable name, similar to the StdStorage UX in forge-std.
• Using variable names would likely require knowing the storage layout (and having forge compile with that enabled by default). But maybe that's ok / a better route, because it would also unlock the ability to add cheats to make it easier to write to packed slots (ref Cannot find the storage slot for a public string variable #3869 (comment)).

[clouds56]

Having to use slot numbers isn't ideal since it's not very user friendly, it would be nice if you can reference mappings by their variable name, similar to the StdStorage UX in forge-std.

I agree with this is not the best UX, but I think it is also important that we could have simple and powerful API in vm.
Since not all mappings has a "name", think of a nest mapping

mapping(uint => mapping(uint => uint)) d;

how to name d[1] which is also a map?

Shall we let StdStorage handle the name resolving things (maybe a separate PR) while let vm simple?
A possible StdStorage API could be like:

// d[1][2] = 3
// `read_mapping` would set storage dynamic type to mapping
StdStorage root = stdstore.target(address(test)).sig("d()").read_mapping();
// `getLength` would call `getMappingLength` on correspond slot
for (uint i = 0; i < root.getLength(); i++) {
  // `element` would call `getMappingSlotAt`
  StdStorage sub = root.element(i).read_mapping();
  for (uint j = 0; j < sub.getLength(); j++) {
    uint value = root.element(j).read_u256();
  }
}

inner StdStorage, these vm API are actually called

            getMappingLength(address,bytes32)
            getMappingSlotAt(address,bytes32,uint256)

[clouds56]

I implemented root() and parent() in StdStorage.
Now we could use

(uint256 slot, bytes32 key) = stdstore.target(address(test)).sig(test.map_addr.selector).with_key(address(this)).parent();
// or
stdstore.target(address(test)).sig(test.map_addr.selector).with_key(address(this)).root();

to get root slot.
Note: you have to pass right number of .with_key, while the key doesn't have to exists. See foundry-rs/forge-std#343 for more details

[mattsse]

code lgtm,

pending @mds1

[mds1]

Sorry for the slow response here, busy week. Lots of questions below.

mapping(uint => mapping(uint => uint)) d;

how to name d[1] which is also a map?

You could do it just like that, that seems like a good convention, i.e. d[0] would be the first mapping key, d[1] the second key, and if there's only a single key then d and d[0] would equivalent.

Shall we let StdStorage handle the name resolving things (maybe a separate PR) while let vm simple?

I do think we should make as much of this native to forge as we can, per #3782, so if we are able to avoid needing a forge-std PR that would be preferable. Rust is a lot more flexible than solidity too, so handling things here seems simpler than having to handle them in solidity. Is there a reason you prefer to split things up into forge-std?

Also can you provide some examples of what it would look like to use the full set of cheats + forge-std to get an array of all mapping keys for a regular mapping and nested mapping? Want to make sure I fully understand the UX here, there's a lot of cheats and I'm not sure I fully understand all the nomenclature yet, e.g. parent, root, etc.

One thing in particular is, instead of a cheat to return the length, should we additionally have a family of cheats like the below? It feels like that would be better UX since I imagine getting the full array is the most common use case, and of course you can always easily get the length from there.

• getMappingKeysUint(...) returns (uint[] memory)
• getMappingKeysAddress(...) returns (address[] memory)
• getMappingKeysBytes(...) returns (bytes[] memory) -- for custom types that aren't bool/uint/int/bytes32/bytes/string/address

It's probably keep to keep the length + index versions for better performance in cases where the arrays get really big.

Should there be a stopMappingRecording() cheat also? Wondering if recording all this can get costly / come with a performance penalty, so you might want the ability to turn it on/off.

Lastly, how do the getters behave in terms of their buffers? For example, recordLogs starts recording logs, then getRecordedLogs returns all recorded logs then clears it's buffer. So an immediate subsequent call to getRecordedLogs returns nothing. We have 4 getters here, so I assume we can't immediately clear the buffer as soon as one is called, and that the lengths only grow?

[clouds56]

I'm not familiar with UX things, the new adding

getMappingKeysUint(...) returns (uint[] memory)
getMappingKeysAddress(...) returns (address[] memory)
getMappingKeysBytes(...) returns (bytes[] memory)

could all built on top of current APIs.

I do think we should make as much of this native to forge as we can, per #3782, so if we are able to avoid needing a forge-std PR that would be preferable.
Is there a reason you prefer to split things up into forge-std?

I like the idea to get rid of forge-std, while if i read it correctly, the #3782 is just a draft RFC, it even hasn't reached the final design. We could not let existing PR to follow a standard not exists.
The reason I prefer split is now things are handled in this way, I'm just making consistence to existing code base.

Should there be a stopMappingRecording() cheat also?

I agree with we should add stopMappingRecording, given we have record(), recordLogs(), I think you should add stopRecord() and stopRecordLogs() as well.

I think the UX of contribution is some how not the best for now? I think we could:

1. let the basic API merged if there's no bug (since any other API could be built upon current API, do you agree with this @mattsse?)
2. gives a final design of implementable API (no matter forge-std or vm itself), and maybe I could do that in a separate PR if I have time.
3. migrate current code to "the better UX" you mean ASAP, so following contributor could learn the best practice by reading the code.

[clouds56]

As a side note for stdstore:
Why couldn't we get slot from layout, forge inspect is layout aware, why can't stdstore get this information?

$ forge inspect src/Greeter.sol:Greeter storage --pretty
| Name     | Type                                           | Slot | Offset | Bytes | Contract                |
|----------|------------------------------------------------|------|--------|-------|-------------------------|
| greeting | string                                         | 0    | 0      | 32    | src/Greeter.sol:Greeter |
| owner    | address                                        | 1    | 0      | 20    | src/Greeter.sol:Greeter |
| data     | mapping(address => uint256)                    | 2    | 0      | 32    | src/Greeter.sol:Greeter |
| data2    | mapping(address => mapping(address => int256)) | 3    | 0      | 32    | src/Greeter.sol:Greeter |

[flyq]

According to my experience here, I think the interface design is quite ergonomic.

Do you have any comments and suggestions for subsequent updates and modifications? @mds1

[mattsse]

I agree with we should add stopMappingRecording, given we have record(), recordLogs(), I think you should add stopRecord() and stopRecordLogs() as well.

this sounds reasonable to me

let the basic API merged if there's no bug (since any other API could be built upon current API, do you agree with this

defer to @mds1 here

[mds1]

let the basic API merged if there's no bug (since any other API could be built upon current API, do you agree with this @mattsse?)

Ok, I'm good with that. I believe you're right we should be able to build the UX improvements on top of this base, and this is better than the existing solution of having to manually implement/track everything anyway

As a side note for stdstore: Why couldn't we get slot from layout, forge inspect is layout aware, why can't stdstore get this information?

By default solc/forge does not include the storage layout in the build artifacts. So we'd either have to:

• Always enable storage layout as an output by default. This might increase compilation times (or maybe not, and it's roughly free), so we'd have to check that first
• Revert if the artifact doesn't have storage layout and explain it needs to be enabled

But otherwise you're right that's were we'd get the info from. forge inspect just recompiles with the storage layout output enabled to get the data

According to my experience here, I think the interface design is quite ergonomic.

This is a good example, thank you.

Do you have any comments and suggestions for subsequent updates and modifications

Here's my suggested plan. First for this PR:

• Merge it (pending any comments @mattsse might have on code)
• Open a PR to the book repo for docs
• Open a PR to forge-std to add the cheats to VmSafe in Vm.sol

Then create a separate issue to tract UX improvements using mapping names, where we should:

• Take a few repos that use forge (should do seaport with via-ir, forge-std, maple labs, uniswaps universal router and permit2 are all good options). Use hyperfine to benchmark the difference in compilation with and without storage layout as an output. It's likely solc already has the data so it's cheap to spit out, or it's possible it involves some extra work with nontrivial added time to compilation
• If it's cheap, let's include it by default in forge for better UX. Having access to storage layout will help a lot with these cheats + writing to packed slots. If it's not cheap, the rest of the steps are the same, but if the new UX-friendlier cheats are used without storage layout they just revert and explain why
• Add cheats that let you specify vars/slots by name instead of index

Thanks a lot for this PR! This will be a great improvement, and I appreciate your patience around the UX questions/decisions :)

[bytes032]

i've just stumbled upon this.. absolutely banger feature

[mds1]

Agreed! @clouds56 any chance you could resolve the conflicts here and maybe we can get this merged this week?

[mattsse]

very nice!

[mattsse]

@clouds56 I'm unable to push directly to this branch (perhaps you don't allow changes by maintainers?)

I pushed to #5123 feel free to update this PR with that and we can close #5123

[Evalir]

merged thru #5123