diff --git a/tests/unit/cli/storage_layout/test_storage_layout.py b/tests/unit/cli/storage_layout/test_storage_layout.py index 9724dd723c..ece2743b81 100644 --- a/tests/unit/cli/storage_layout/test_storage_layout.py +++ b/tests/unit/cli/storage_layout/test_storage_layout.py @@ -1,4 +1,21 @@ from vyper.compiler import compile_code +from vyper.evm.opcodes import version_check + + +def _adjust_storage_layout_for_cancun(layout): + def _go(layout): + for _varname, item in layout.items(): + if "slot" in item and isinstance(item["slot"], int): + item["slot"] -= 1 + else: + # recurse to submodule + _go(item) + + if version_check(begin="cancun"): + layout["transient_storage_layout"] = { + "$.nonreentrant_key": layout["storage_layout"].pop("$.nonreentrant_key") + } + _go(layout["storage_layout"]) def test_storage_layout(): @@ -40,13 +57,18 @@ def public_foo3(): out = compile_code(code, output_formats=["layout"]) - assert out["layout"]["storage_layout"] == { - "$.nonreentrant_key": {"slot": 0, "type": "nonreentrant lock"}, - "foo": {"slot": 1, "type": "HashMap[address, uint256]"}, - "arr": {"slot": 2, "type": "DynArray[uint256, 3]"}, - "baz": {"slot": 6, "type": "Bytes[65]"}, - "bar": {"slot": 10, "type": "uint256"}, + expected = { + "storage_layout": { + "$.nonreentrant_key": {"slot": 0, "type": "nonreentrant lock"}, + "foo": {"slot": 1, "type": "HashMap[address, uint256]"}, + "arr": {"slot": 2, "type": "DynArray[uint256, 3]"}, + "baz": {"slot": 6, "type": "Bytes[65]"}, + "bar": {"slot": 10, "type": "uint256"}, + } } + _adjust_storage_layout_for_cancun(expected) + + assert out["layout"] == expected def test_storage_and_immutables_layout(): @@ -71,6 +93,7 @@ def __init__(): "name": {"slot": 1, "type": "String[32]"}, }, } + _adjust_storage_layout_for_cancun(expected_layout) out = compile_code(code, output_formats=["layout"]) assert out["layout"] == expected_layout @@ -120,6 +143,7 @@ def __init__(): "a_library": {"supply": {"slot": 3, "type": "uint256"}}, }, } + _adjust_storage_layout_for_cancun(expected_layout) out = compile_code(code, input_bundle=input_bundle, output_formats=["layout"]) assert out["layout"] == expected_layout @@ -169,6 +193,7 @@ def __init__(): "counter2": {"slot": 3, "type": "uint256"}, }, } + _adjust_storage_layout_for_cancun(expected_layout) out = compile_code(code, input_bundle=input_bundle, output_formats=["layout"]) assert out["layout"] == expected_layout @@ -253,6 +278,7 @@ def bar(): "a_library": {"supply": {"slot": 4, "type": "uint256"}}, }, } + _adjust_storage_layout_for_cancun(expected_layout) out = compile_code(code, input_bundle=input_bundle, output_formats=["layout"]) assert out["layout"] == expected_layout @@ -334,6 +360,7 @@ def foo() -> uint256: "counter2": {"slot": 4, "type": "uint256"}, }, } + _adjust_storage_layout_for_cancun(expected_layout) out = compile_code(code, input_bundle=input_bundle, output_formats=["layout"]) assert out["layout"] == expected_layout