refactor[tool]: refactor storage layout export (#3789)

refactor storage layout allocator. separate concerns of allocating the storage layout and exporting the storage layout into separate functions. this is intended to make it easier to add features to the storage layout export in the future fix several bugs in storage layout overrides, including: - handle stateful modules - add a sanity check that the override file roundtrips - ignore non-storage variables in override files - set nonreentrant lock properly for all functions instead of panicking misc: add `n_slots` to each storage layout item in the export
vyperlang · May 28, 2024 · 96a8384 · 96a8384
1 parent 608feda
commit 96a8384
Show file tree

Hide file tree

Showing 7 changed files with 499 additions and 152 deletions.
diff --git a/tests/unit/cli/storage_layout/__init__.py b/tests/unit/cli/storage_layout/__init__.py
diff --git a/tests/unit/cli/storage_layout/test_storage_layout.py b/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