feat: Storage Layout Overrides

tests/cli/outputs/test_storage_layout_overrides.py

@@ -0,0 +1,114 @@
+import pytest
+
+from vyper.compiler import compile_code
+from vyper.exceptions import StorageLayoutException
+
+
+def test_storage_layout_overrides():
+    code = """
+a: uint256
+b: uint256"""
+
+    storage_layout_overrides = {
+        "a": {"type": "uint256", "location": "storage", "slot": 1},
+        "b": {"type": "uint256", "location": "storage", "slot": 0},
+    }
+
+    out = compile_code(
+        code, output_formats=["layout"], storage_layout_override=storage_layout_overrides
+    )
+
+    assert out["layout"] == storage_layout_overrides
+
+
+def test_storage_layout_for_more_complex():
+    code = """
+foo: HashMap[address, uint256]
+
+@external
+@nonreentrant("foo")
+def public_foo1():
+    pass
+
+@external
+@nonreentrant("foo")
+def public_foo2():
+    pass
+
+
+@internal
+@nonreentrant("bar")
+def _bar():
+    pass
+
+# mix it up a little
+baz: Bytes[65]
+bar: uint256
+
+@external
+@nonreentrant("bar")
+def public_bar():
+    pass
+
+@external
+@nonreentrant("foo")
+def public_foo3():
+    pass
+    """
+
+    storage_layout_override = {
+        "nonreentrant.foo": {"type": "nonreentrant lock", "location": "storage", "slot": 8},
+        "nonreentrant.bar": {"type": "nonreentrant lock", "location": "storage", "slot": 7},
+        "foo": {
+            "type": "HashMap[address, uint256]",
+            "location": "storage",
+            "slot": 1,
+        },
+        "baz": {"type": "Bytes[65]", "location": "storage", "slot": 2},
+        "bar": {"type": "uint256", "location": "storage", "slot": 6},
+    }
+
+    out = compile_code(
+        code, output_formats=["layout"], storage_layout_override=storage_layout_override
+    )
+
+    assert out["layout"] == storage_layout_override
+
+
+def test_simple_collision():
+    code = """
+name: public(String[64])
+symbol: public(String[32])"""
+
+    storage_layout_override = {
+        "name": {"location": "storage", "slot": 0, "type": "String[64]"},
+        "symbol": {"location": "storage", "slot": 1, "type": "String[32]"},
+    }
+
+    with pytest.raises(
+        StorageLayoutException,
+        match="Storage collision! Tried to assign 'symbol' to slot 1"
+        " but it has already been reserved by 'name'",
+    ):
+        compile_code(
+            code, output_formats=["layout"], storage_layout_override=storage_layout_override
+        )
+
+
+def test_incomplete_overrides():
+    code = """
+name: public(String[64])
+symbol: public(String[32])"""
+
+    storage_layout_override = {
+        "name": {"location": "storage", "slot": 0, "type": "String[64]"},
+    }
+
+    with pytest.raises(
+        StorageLayoutException,
+        match="Could not find storage_slot for symbol. "
+        "Have you used the correct storage layout file?",
+    ):
+        compile_code(
+            code, output_formats=["layout"], storage_layout_override=storage_layout_override
+        )

vyper/cli/vyper_compile.py

@@ -99,6 +99,12 @@ def _parse_args(argv):
         default="bytecode",
         dest="format",
     )
+    parser.add_argument(
+        "--storage-layout-file",
+        help="Override storage slots provided by compiler",
+        dest="storage_layout",
+        nargs="+",
+    )
     parser.add_argument(
         "--evm-version",
         help=f"Select desired EVM version (default {DEFAULT_EVM_VERSION})",
@@ -156,6 +162,7 @@ def _parse_args(argv):
         args.show_gas_estimates,
         args.evm_version,
         args.no_optimize,
+        args.storage_layout,
     )
 
     if args.output_path:
@@ -227,6 +234,7 @@ def compile_files(
     show_gas_estimates: bool = False,
     evm_version: str = DEFAULT_EVM_VERSION,
     no_optimize: bool = False,
+    storage_layout: Iterable[str] = None,
 ) -> OrderedDict:
 
     root_path = Path(root_folder).resolve()
@@ -245,6 +253,13 @@ def compile_files(
             # https://bugs.python.org/issue35107
             contract_sources[file_str] = fh.read() + "\n"
 
+    storage_layouts = OrderedDict()
+    if storage_layout:
+        for storage_file_name, contract_name in zip(storage_layout, contract_sources.keys()):
+            storage_file_path = Path(storage_file_name)
+            with storage_file_path.open() as sfh:
+                storage_layouts[contract_name] = json.load(sfh)
+
     show_version = False
     if "combined_json" in output_formats:
         if len(output_formats) > 1:
@@ -262,6 +277,7 @@ def compile_files(
         interface_codes=get_interface_codes(root_path, contract_sources),
         evm_version=evm_version,
         no_optimize=no_optimize,
+        storage_layouts=storage_layouts,
         show_gas_estimates=show_gas_estimates,
     )
     if show_version:

vyper/compiler/__init__.py

@@ -1,10 +1,18 @@
 from collections import OrderedDict
-from typing import Any, Callable, Optional, Sequence, Union
+from typing import Any, Callable, Dict, Optional, Sequence, Union
 
 from vyper.compiler import output
 from vyper.compiler.phases import CompilerData
 from vyper.evm.opcodes import DEFAULT_EVM_VERSION, evm_wrapper
-from vyper.typing import ContractCodes, InterfaceDict, InterfaceImports, OutputDict, OutputFormats
+from vyper.typing import (
+    ContractCodes,
+    ContractPath,
+    InterfaceDict,
+    InterfaceImports,
+    OutputDict,
+    OutputFormats,
+    StorageLayout,
+)
 
 OUTPUT_FORMATS = {
     # requires vyper_module
@@ -39,6 +47,7 @@ def compile_codes(
     interface_codes: Union[InterfaceDict, InterfaceImports, None] = None,
     initial_id: int = 0,
     no_optimize: bool = False,
+    storage_layouts: Dict[ContractPath, StorageLayout] = None,
     show_gas_estimates: bool = False,
 ) -> OrderedDict:
     """
@@ -89,6 +98,10 @@ def compile_codes(
     for source_id, contract_name in enumerate(sorted(contract_sources), start=initial_id):
         source_code = contract_sources[contract_name]
         interfaces: Any = interface_codes
+        storage_layout_override = None
+        if storage_layouts and contract_name in storage_layouts:
+            storage_layout_override = storage_layouts[contract_name]
+
         if (
             isinstance(interfaces, dict)
             and contract_name in interfaces
@@ -97,12 +110,7 @@ def compile_codes(
             interfaces = interfaces[contract_name]
 
         compiler_data = CompilerData(
-            source_code,
-            contract_name,
-            interfaces,
-            source_id,
-            no_optimize,
-            show_gas_estimates
+            source_code, contract_name, interfaces, source_id, no_optimize, storage_layout_override, show_gas_estimates
         )
         for output_format in output_formats[contract_name]:
             if output_format not in OUTPUT_FORMATS:
@@ -128,6 +136,7 @@ def compile_code(
     interface_codes: Optional[InterfaceImports] = None,
     evm_version: str = DEFAULT_EVM_VERSION,
     no_optimize: bool = False,
+    storage_layout_override: StorageLayout = None,
     show_gas_estimates: bool = False,
 ) -> dict:
     """
@@ -160,12 +169,14 @@ def compile_code(
     """
 
     contract_sources = {UNKNOWN_CONTRACT_NAME: contract_source}
+    storage_layouts = {UNKNOWN_CONTRACT_NAME: storage_layout_override}
 
     return compile_codes(
         contract_sources,
         output_formats,
         interface_codes=interface_codes,
         evm_version=evm_version,
         no_optimize=no_optimize,
+        storage_layouts=storage_layouts,
         show_gas_estimates=show_gas_estimates,
     )[UNKNOWN_CONTRACT_NAME]

vyper/compiler/phases.py

@@ -48,6 +48,7 @@ def __init__(
         interface_codes: Optional[InterfaceImports] = None,
         source_id: int = 0,
         no_optimize: bool = False,
+        storage_layout: StorageLayout = None,
         show_gas_estimates: bool = False,
     ) -> None:
         """
@@ -75,6 +76,7 @@ def __init__(
         self.interface_codes = interface_codes
         self.source_id = source_id
         self.no_optimize = no_optimize
+        self.storage_layout_override = storage_layout
         self.show_gas_estimates = show_gas_estimates
 
     @property
@@ -88,7 +90,7 @@ def vyper_module(self) -> vy_ast.Module:
     def vyper_module_folded(self) -> vy_ast.Module:
         if not hasattr(self, "_vyper_module_folded"):
             self._vyper_module_folded, self._storage_layout = generate_folded_ast(
-                self.vyper_module, self.interface_codes
+                self.vyper_module, self.interface_codes, self.storage_layout_override
             )
 
         return self._vyper_module_folded
@@ -97,7 +99,7 @@ def vyper_module_folded(self) -> vy_ast.Module:
     def storage_layout(self) -> StorageLayout:
         if not hasattr(self, "_storage_layout"):
             self._vyper_module_folded, self._storage_layout = generate_folded_ast(
-                self.vyper_module, self.interface_codes
+                self.vyper_module, self.interface_codes, self.storage_layout_override
             )
 
         return self._storage_layout
@@ -174,7 +176,9 @@ def generate_ast(source_code: str, source_id: int, contract_name: str) -> vy_ast
 
 
 def generate_folded_ast(
-    vyper_module: vy_ast.Module, interface_codes: Optional[InterfaceImports]
+    vyper_module: vy_ast.Module,
+    interface_codes: Optional[InterfaceImports],
+    storage_layout_overrides: StorageLayout = None,
 ) -> Tuple[vy_ast.Module, StorageLayout]:
     """
     Perform constant folding operations on the Vyper AST.
@@ -197,7 +201,7 @@ def generate_folded_ast(
     vy_ast.folding.fold(vyper_module_folded)
     validate_semantics(vyper_module_folded, interface_codes)
     vy_ast.expansion.expand_annotated_ast(vyper_module_folded)
-    symbol_tables = set_data_positions(vyper_module_folded)
+    symbol_tables = set_data_positions(vyper_module_folded, storage_layout_overrides)
 
     return vyper_module_folded, symbol_tables
 

vyper/exceptions.py

@@ -255,6 +255,10 @@ class EvmVersionException(VyperException):
     """Invalid action for the active EVM ruleset."""
 
 
+class StorageLayoutException(VyperException):
+    """Invalid slot for the storage layout overrides"""
+
+
 class JSONError(Exception):
 
     """Invalid compiler input JSON."""