refactor[tool]: refactor storage layout export

tests/unit/cli/storage_layout/test_storage_layout.py

@@ -1,21 +1,6 @@
 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"])
+from .utils import adjust_storage_layout_for_cancun
 
 
 def test_storage_layout():
@@ -55,19 +40,18 @@ def public_foo3():
     pass
     """
 
-    out = compile_code(code, output_formats=["layout"])
-
     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"},
+            "$.nonreentrant_key": {"slot": 0, "type": "nonreentrant lock", "n_slots": 1},
+            "foo": {"slot": 1, "type": "HashMap[address, uint256]", "n_slots": 1},
+            "arr": {"slot": 2, "type": "DynArray[uint256, 3]", "n_slots": 4},
+            "baz": {"slot": 6, "type": "Bytes[65]", "n_slots": 4},
+            "bar": {"slot": 10, "type": "uint256", "n_slots": 1},
         }
     }
-    _adjust_storage_layout_for_cancun(expected)
+    adjust_storage_layout_for_cancun(expected)
 
+    out = compile_code(code, output_formats=["layout"])
     assert out["layout"] == expected
 
 
@@ -88,12 +72,9 @@ def __init__():
             "SYMBOL": {"length": 64, "offset": 0, "type": "String[32]"},
             "DECIMALS": {"length": 32, "offset": 64, "type": "uint8"},
         },
-        "storage_layout": {
-            "$.nonreentrant_key": {"slot": 0, "type": "nonreentrant lock"},
-            "name": {"slot": 1, "type": "String[32]"},
-        },
+        "storage_layout": {"name": {"slot": 1, "type": "String[32]", "n_slots": 2}},
     }
-    _adjust_storage_layout_for_cancun(expected_layout)
+    adjust_storage_layout_for_cancun(expected_layout)
 
     out = compile_code(code, output_formats=["layout"])
     assert out["layout"] == expected_layout
@@ -137,13 +118,12 @@ def __init__():
             },
         },
         "storage_layout": {
-            "$.nonreentrant_key": {"slot": 0, "type": "nonreentrant lock"},
-            "counter": {"slot": 1, "type": "uint256"},
-            "counter2": {"slot": 2, "type": "uint256"},
-            "a_library": {"supply": {"slot": 3, "type": "uint256"}},
+            "counter": {"slot": 1, "type": "uint256", "n_slots": 1},
+            "counter2": {"slot": 2, "type": "uint256", "n_slots": 1},
+            "a_library": {"supply": {"slot": 3, "type": "uint256", "n_slots": 1}},
         },
     }
-    _adjust_storage_layout_for_cancun(expected_layout)
+    adjust_storage_layout_for_cancun(expected_layout)
 
     out = compile_code(code, input_bundle=input_bundle, output_formats=["layout"])
     assert out["layout"] == expected_layout
@@ -187,13 +167,12 @@ def __init__():
             },
         },
         "storage_layout": {
-            "$.nonreentrant_key": {"slot": 0, "type": "nonreentrant lock"},
-            "counter": {"slot": 1, "type": "uint256"},
-            "a_library": {"supply": {"slot": 2, "type": "uint256"}},
-            "counter2": {"slot": 3, "type": "uint256"},
+            "counter": {"slot": 1, "type": "uint256", "n_slots": 1},
+            "a_library": {"supply": {"slot": 2, "type": "uint256", "n_slots": 1}},
+            "counter2": {"slot": 3, "type": "uint256", "n_slots": 1},
         },
     }
-    _adjust_storage_layout_for_cancun(expected_layout)
+    adjust_storage_layout_for_cancun(expected_layout)
 
     out = compile_code(code, input_bundle=input_bundle, output_formats=["layout"])
     assert out["layout"] == expected_layout
@@ -271,14 +250,14 @@ def bar():
             },
         },
         "storage_layout": {
-            "$.nonreentrant_key": {"slot": 0, "type": "nonreentrant lock"},
-            "counter": {"slot": 1, "type": "uint256"},
-            "lib2": {"storage_variable": {"slot": 2, "type": "uint256"}},
-            "counter2": {"slot": 3, "type": "uint256"},
-            "a_library": {"supply": {"slot": 4, "type": "uint256"}},
+            "$.nonreentrant_key": {"slot": 0, "type": "nonreentrant lock", "n_slots": 1},
+            "counter": {"slot": 1, "type": "uint256", "n_slots": 1},
+            "lib2": {"storage_variable": {"slot": 2, "type": "uint256", "n_slots": 1}},
+            "counter2": {"slot": 3, "type": "uint256", "n_slots": 1},
+            "a_library": {"supply": {"slot": 4, "type": "uint256", "n_slots": 1}},
         },
     }
-    _adjust_storage_layout_for_cancun(expected_layout)
+    adjust_storage_layout_for_cancun(expected_layout)
 
     out = compile_code(code, input_bundle=input_bundle, output_formats=["layout"])
     assert out["layout"] == expected_layout
@@ -351,16 +330,15 @@ def foo() -> uint256:
             },
         },
         "storage_layout": {
-            "$.nonreentrant_key": {"slot": 0, "type": "nonreentrant lock"},
-            "counter": {"slot": 1, "type": "uint256"},
+            "counter": {"slot": 1, "type": "uint256", "n_slots": 1},
             "lib2": {
-                "lib1": {"supply": {"slot": 2, "type": "uint256"}},
-                "storage_variable": {"slot": 3, "type": "uint256"},
+                "lib1": {"supply": {"slot": 2, "type": "uint256", "n_slots": 1}},
+                "storage_variable": {"slot": 3, "type": "uint256", "n_slots": 1},
             },
-            "counter2": {"slot": 4, "type": "uint256"},
+            "counter2": {"slot": 4, "type": "uint256", "n_slots": 1},
         },
     }
-    _adjust_storage_layout_for_cancun(expected_layout)
+    adjust_storage_layout_for_cancun(expected_layout)
 
     out = compile_code(code, input_bundle=input_bundle, output_formats=["layout"])
     assert out["layout"] == expected_layout

tests/unit/cli/storage_layout/test_storage_layout_overrides.py

@@ -5,18 +5,20 @@
 from vyper.compiler import compile_code
 from vyper.exceptions import StorageLayoutException
 
+from .utils import adjust_storage_layout_for_cancun
+
 
 def test_storage_layout_overrides():
     code = """
 a: uint256
 b: uint256"""
 
     storage_layout_overrides = {
-        "a": {"type": "uint256", "slot": 1},
-        "b": {"type": "uint256", "slot": 0},
+        "a": {"type": "uint256", "slot": 1, "n_slots": 1},
+        "b": {"type": "uint256", "slot": 0, "n_slots": 1},
     }
 
-    expected_output = {"storage_layout": storage_layout_overrides, "code_layout": {}}
+    expected_output = {"storage_layout": storage_layout_overrides}
 
     out = compile_code(
         code, output_formats=["layout"], storage_layout_override=storage_layout_overrides
@@ -61,18 +63,20 @@ def public_foo3():
     """
 
     storage_layout_override = {
-        "$.nonreentrant_key": {"type": "nonreentrant lock", "slot": 8},
-        "foo": {"type": "HashMap[address, uint256]", "slot": 1},
-        "baz": {"type": "Bytes[65]", "slot": 2},
-        "bar": {"type": "uint256", "slot": 6},
+        "$.nonreentrant_key": {"type": "nonreentrant lock", "slot": 8, "n_slots": 1},
+        "foo": {"type": "HashMap[address, uint256]", "slot": 1, "n_slots": 1},
+        "baz": {"type": "Bytes[65]", "slot": 2, "n_slots": 4},
+        "bar": {"type": "uint256", "slot": 6, "n_slots": 1},
     }
 
-    expected_output = {"storage_layout": storage_layout_override, "code_layout": {}}
+    expected_output = {"storage_layout": storage_layout_override}
 
     out = compile_code(
         code, output_formats=["layout"], storage_layout_override=storage_layout_override
     )
 
+    adjust_storage_layout_for_cancun(expected_output, do_adjust_slots=False)
+
     assert out["layout"] == expected_output
 
 

tests/unit/cli/storage_layout/utils.py

@@ -0,0 +1,18 @@
+from vyper.evm.opcodes import version_check
+
+
+def adjust_storage_layout_for_cancun(layout, do_adjust_slots=True):
+    def _go(layout):
+        for _varname, item in layout.items():
+            if "slot" in item and isinstance(item["slot"], int):
+                if do_adjust_slots:
+                    item["slot"] -= 1
+            else:
+                # recurse to submodule
+                _go(item)
+
+    if version_check(begin="cancun"):
+        nonreentrant = layout["storage_layout"].pop("$.nonreentrant_key", None)
+        if nonreentrant is not None:
+            layout["transient_storage_layout"] = {"$.nonreentrant_key": nonreentrant}
+        _go(layout["storage_layout"])

vyper/compiler/phases.py

@@ -13,6 +13,7 @@
 from vyper.exceptions import StructureException
 from vyper.ir import compile_ir, optimizer
 from vyper.semantics import analyze_module, set_data_positions, validate_compilation_target
+from vyper.semantics.analysis.data_positions import generate_layout_export
 from vyper.semantics.types.function import ContractFunctionT
 from vyper.semantics.types.module import ModuleT
 from vyper.typing import StorageLayout
@@ -167,7 +168,9 @@ def compilation_target(self):
     @cached_property
     def storage_layout(self) -> StorageLayout:
         module_ast = self.compilation_target
-        return set_data_positions(module_ast, self.storage_layout_override)
+        set_data_positions(module_ast, self.storage_layout_override)
+
+        return generate_layout_export(module_ast)
 
     @property
     def global_ctx(self) -> ModuleT: