Repository: vyperlang/vyper

Issue:
tool: allow storage layout overrides in standard-json and vyz files
we allow storage layout overrides from the CLI, however, not from the other compiler entry points. we need to add those so verifiers can verify storage layout overrides.

Discussion:
charles-cooper: @tserg looks like the latest merge from master broke some tests - could you take a look?

Patch:
diff --git a/docs/compiling-a-contract.rst b/docs/compiling-a-contract.rst
index 7132cff5..c839e1e8 100644
--- a/docs/compiling-a-contract.rst
+++ b/docs/compiling-a-contract.rst
@@ -106,7 +106,7 @@ Online Compilers
 Try VyperLang!
 -----------------
 
-`Try VyperLang! <https://try.vyperlang.org>`_ is a JupterHub instance hosted by the Vyper team as a sandbox for developing and testing contracts in Vyper. It requires github for login, and supports deployment via the browser.
+`Try VyperLang! <https://try.vyperlang.org>`_ is a JupyterHub instance hosted by the Vyper team as a sandbox for developing and testing contracts in Vyper. It requires github for login, and supports deployment via the browser.
 
 Remix IDE
 ---------
@@ -203,7 +203,7 @@ The following is a list of supported EVM versions, and changes in the compiler i
 Integrity Hash
 ==============
 
-To help tooling detect whether two builds are the same, Vyper provides the ``-f integrity`` output, which outputs the integrity hash of a contract. The integrity hash is recursively defined as the sha256 of the source code with the integrity hashes of its dependencies (imports).
+To help tooling detect whether two builds are the same, Vyper provides the ``-f integrity`` output, which outputs the integrity hash of a contract. The integrity hash is recursively defined as the sha256 of the source code with the integrity hashes of its dependencies (imports) and storage layout overrides (if provided).
 
 .. _vyper-archives:
 
@@ -219,8 +219,9 @@ A Vyper archive is a compileable bundle of input sources and settings. Technical
     ├── compilation_targets
     ├── compiler_version
     ├── integrity
+    ├── settings.json
     ├── searchpaths
-    └── settings.json
+    └── storage_layout.json [OPTIONAL]
 
 * ``cli_settings.txt`` is a text representation of the settings that were used on the compilation run that generated this archive.
 * ``compilation_targets`` is a newline separated list of compilation targets. Currently only one compilation is supported
@@ -228,6 +229,7 @@ A Vyper archive is a compileable bundle of input sources and settings. Technical
 * ``integrity`` is the :ref:`integrity hash <integrity-hash>` of the input contract
 * ``searchpaths`` is a newline-separated list of the search paths used on this compilation run
 * ``settings.json`` is a json representation of the settings used on this compilation run. It is 1:1 with ``cli_settings.txt``, but both are provided as they are convenient for different workflows (typically, manually vs automated).
+* ``storage_layout.json`` is a json representation of the storage layout overrides to be used on this compilation run. It is optional.
 
 A Vyper archive file can be produced by requesting the ``-f archive`` output format. The compiler can also produce the archive in base64 encoded form using the ``--base64`` flag. The Vyper compiler can accept both ``.vyz`` and base64-encoded Vyper archives directly as input.
 
@@ -281,6 +283,14 @@ The following example describes the expected input format of ``vyper-json``. (Co
             }
         },
         // Optional
+        // Storage layout overrides for the contracts that are compiled
+        "storage_layout_overrides": {
+            "contracts/foo.vy": {
+                "a": {"type": "uint256", "slot": 1, "n_slots": 1},
+                "b": {"type": "uint256", "slot": 0, "n_slots": 1},
+            }
+        },
+        // Optional
         "settings": {
             "evmVersion": "cancun",  // EVM version to compile for. Can be london, paris, shanghai or cancun (default).
             // optional, optimization mode
@@ -364,6 +374,13 @@ The following example describes the output format of ``vyper-json``. Comments ar
             "formattedMessage": "line 5:11 Unsupported type conversion: int128 to bool"
             }
         ],
+        // Optional: not present if there are no storage layout overrides
+        "storage_layout_overrides": {
+            "contracts/foo.vy": {
+                "a": {"type": "uint256", "slot": 1, "n_slots": 1},
+                "b": {"type": "uint256", "slot": 0, "n_slots": 1},
+            }
+        },
         // This contains the file-level outputs. Can be limited/filtered by the outputSelection settings.
         "sources": {
             "source_file.vy": {
diff --git a/vyper/cli/compile_archive.py b/vyper/cli/compile_archive.py
index c6d07de9..d1dd2588 100644
--- a/vyper/cli/compile_archive.py
+++ b/vyper/cli/compile_archive.py
@@ -45,6 +45,11 @@ def compiler_data_from_zip(file_name, settings, no_bytecode_metadata):
     fcontents = archive.read("MANIFEST/compilation_targets").decode("utf-8")
     compilation_targets = fcontents.splitlines()
 
+    storage_layout_path = "MANIFEST/storage_layout.json"
+    storage_layout = None
+    if storage_layout_path in archive.namelist():
+        storage_layout = json.loads(archive.read(storage_layout_path).decode("utf-8"))
+
     if len(compilation_targets) != 1:
         raise BadArchive("Multiple compilation targets not supported!")
 
@@ -68,6 +73,7 @@ def compiler_data_from_zip(file_name, settings, no_bytecode_metadata):
     return CompilerData(
         file,
         input_bundle=input_bundle,
+        storage_layout=storage_layout,
         integrity_sum=integrity,
         settings=settings,
         no_bytecode_metadata=no_bytecode_metadata,
diff --git a/vyper/cli/vyper_json.py b/vyper/cli/vyper_json.py
index e7704b93..5f632f41 100755
--- a/vyper/cli/vyper_json.py
+++ b/vyper/cli/vyper_json.py
@@ -12,6 +12,7 @@ from vyper.compiler.input_bundle import FileInput, JSONInputBundle
 from vyper.compiler.settings import OptimizationLevel, Settings
 from vyper.evm.opcodes import EVM_VERSIONS
 from vyper.exceptions import JSONError
+from vyper.typing import StorageLayout
 from vyper.utils import OrderedSet, keccak256
 
 TRANSLATE_MAP = {
@@ -206,6 +207,19 @@ def get_inputs(input_dict: dict) -> dict[PurePath, Any]:
     return ret
 
 
+def get_storage_layout_overrides(input_dict: dict) -> dict[PurePath, StorageLayout]:
+    storage_layout_overrides: dict[PurePath, StorageLayout] = {}
+
+    for path, value in input_dict.get("storage_layout_overrides", {}).items():
+        if path not in input_dict["sources"]:
+            raise JSONError(f"unknown target for storage layout override: {path}")
+
+        path = PurePath(path)
+        storage_layout_overrides[path] = value
+
+    return storage_layout_overrides
+
+
 # get unique output formats for each contract, given the input_dict
 # NOTE: would maybe be nice to raise on duplicated output formats
 def get_output_formats(input_dict: dict) -> dict[PurePath, list[str]]:
@@ -299,6 +313,7 @@ def compile_from_input_dict(
     integrity = input_dict.get("integrity")
 
     sources = get_inputs(input_dict)
+    storage_layout_overrides = get_storage_layout_overrides(input_dict)
     output_formats = get_output_formats(input_dict)
     compilation_targets = list(output_formats.keys())
     search_paths = get_search_paths(input_dict)
@@ -308,6 +323,7 @@ def compile_from_input_dict(
     res, warnings_dict = {}, {}
     warnings.simplefilter("always")
     for contract_path in compilation_targets:
+        storage_layout_override = storage_layout_overrides.get(contract_path)
         with warnings.catch_warnings(record=True) as caught_warnings:
             try:
                 # use load_file to get a unique source_id
@@ -317,6 +333,7 @@ def compile_from_input_dict(
                     file,
                     input_bundle=input_bundle,
                     output_formats=output_formats[contract_path],
+                    storage_layout_override=storage_layout_override,
                     integrity_sum=integrity,
                     settings=settings,
                     no_bytecode_metadata=no_bytecode_metadata,
@@ -356,6 +373,9 @@ def format_to_output_dict(compiler_data: dict) -> dict:
             if key in data:
                 output_contracts[key] = data[key]
 
+        if "layout" in data:
+            output_contracts["layout"] = data["layout"]
+
         if "method_identifiers" in data:
             output_contracts["evm"] = {"methodIdentifiers": data["method_identifiers"]}
 
diff --git a/vyper/compiler/output.py b/vyper/compiler/output.py
index ca951b8e..b6a0e8ac 100644
--- a/vyper/compiler/output.py
+++ b/vyper/compiler/output.py
@@ -102,7 +102,7 @@ def build_archive_b64(compiler_data: CompilerData) -> str:
 
 
 def build_integrity(compiler_data: CompilerData) -> str:
-    return compiler_data.resolved_imports.integrity_sum
+    return compiler_data.integrity_sum
 
 
 def build_external_interface_output(compiler_data: CompilerData) -> str:
diff --git a/vyper/compiler/output_bundle.py b/vyper/compiler/output_bundle.py
index 24a0d070..8af1b722 100644
--- a/vyper/compiler/output_bundle.py
+++ b/vyper/compiler/output_bundle.py
@@ -12,6 +12,7 @@ from vyper.compiler.phases import CompilerData
 from vyper.compiler.settings import Settings
 from vyper.exceptions import CompilerPanic
 from vyper.semantics.analysis.imports import _is_builtin
+from vyper.typing import StorageLayout
 from vyper.utils import get_long_version, safe_relpath
 
 # data structures and routines for constructing "output bundles",
@@ -134,6 +135,11 @@ class OutputBundleWriter:
     def write_sources(self, sources: dict[str, CompilerInput]):
         raise NotImplementedError(f"write_sources: {self.__class__}")
 
+    def write_storage_layout_overrides(
+        self, compilation_target_path: str, storage_layout_override: StorageLayout
+    ):
+        raise NotImplementedError(f"write_storage_layout_overrides: {self.__class__}")
+
     def write_search_paths(self, search_paths: list[str]):
         raise NotImplementedError(f"write_search_paths: {self.__class__}")
 
@@ -158,8 +164,12 @@ class OutputBundleWriter:
         self.write_compilation_target([self.bundle.compilation_target_path])
         self.write_search_paths(self.bundle.used_search_paths)
         self.write_settings(self.compiler_data.original_settings)
-        self.write_integrity(self.compiler_data.resolved_imports.integrity_sum)
+        self.write_integrity(self.compiler_data.integrity_sum)
         self.write_sources(self.bundle.compiler_inputs)
+        if self.compiler_data.storage_layout_override is not None:
+            self.write_storage_layout_overrides(
+                self.bundle.compilation_target_path, self.compiler_data.storage_layout_override
+            )
 
 
 class SolcJSONWriter(OutputBundleWriter):
@@ -175,6 +185,13 @@ class SolcJSONWriter(OutputBundleWriter):
 
         self._output["sources"].update(out)
 
+    def write_storage_layout_overrides(
+        self, compilation_target_path: str, storage_layout_override: StorageLayout
+    ):
+        self._output["storage_layout_overrides"] = {
+            compilation_target_path: storage_layout_override
+        }
+
     def write_search_paths(self, search_paths: list[str]):
         self._output["settings"]["search_paths"] = search_paths
 
@@ -237,6 +254,11 @@ class VyperArchiveWriter(OutputBundleWriter):
         for path, c in sources.items():
             self.archive.writestr(_anonymize(path), c.contents)
 
+    def write_storage_layout_overrides(
+        self, compilation_target_path: str, storage_layout_override: StorageLayout
+    ):
+        self.archive.writestr("MANIFEST/storage_layout.json", json.dumps(storage_layout_override))
+
     def write_search_paths(self, search_paths: list[str]):
         self.archive.writestr("MANIFEST/searchpaths", "\n".join(search_paths))
 
diff --git a/vyper/compiler/phases.py b/vyper/compiler/phases.py
index e6cb1c58..17812ee5 100644
--- a/vyper/compiler/phases.py
+++ b/vyper/compiler/phases.py
@@ -1,4 +1,5 @@
 import copy
+import json
 import warnings
 from functools import cached_property
 from pathlib import Path, PurePath
@@ -19,7 +20,7 @@ from vyper.semantics.analysis.imports import resolve_imports
 from vyper.semantics.types.function import ContractFunctionT
 from vyper.semantics.types.module import ModuleT
 from vyper.typing import StorageLayout
-from vyper.utils import ERC5202_PREFIX, vyper_warn
+from vyper.utils import ERC5202_PREFIX, sha256sum, vyper_warn
 from vyper.venom import generate_assembly_experimental, generate_ir
 
 DEFAULT_CONTRACT_PATH = PurePath("VyperContract.vy")
@@ -151,29 +152,41 @@ class CompilerData:
         _, ast = self._generate_ast
         return ast
 
+    def _compute_integrity_sum(self, imports_integrity_sum: str) -> str:
+        if self.storage_layout_override is not None:
+            layout_sum = sha256sum(json.dumps(self.storage_layout_override))
+            return sha256sum(layout_sum + imports_integrity_sum)
+        return imports_integrity_sum
+
     @cached_property
     def _resolve_imports(self):
         # deepcopy so as to not interfere with `-f ast` output
         vyper_module = copy.deepcopy(self.vyper_module)
         with self.input_bundle.search_path(Path(vyper_module.resolved_path).parent):
-            return vyper_module, resolve_imports(vyper_module, self.input_bundle)
+            imports = resolve_imports(vyper_module, self.input_bundle)
 
-    @cached_property
-    def resolved_imports(self):
-        imports = self._resolve_imports[1]
+        # check integrity sum
+        integrity_sum = self._compute_integrity_sum(imports._integrity_sum)
 
         expected = self.expected_integrity_sum
-
-        if expected is not None and imports.integrity_sum != expected:
+        if expected is not None and integrity_sum != expected:
             # warn for now. strict/relaxed mode was considered but it costs
             # interface and testing complexity to add another feature flag.
             vyper_warn(
                 f"Mismatched integrity sum! Expected {expected}"
-                f" but got {imports.integrity_sum}."
+                f" but got {integrity_sum}."
                 " (This likely indicates a corrupted archive)"
             )
 
-        return imports
+        return vyper_module, imports, integrity_sum
+
+    @cached_property
+    def integrity_sum(self):
+        return self._resolve_imports[2]
+
+    @cached_property
+    def resolved_imports(self):
+        return self._resolve_imports[1]
 
     @cached_property
     def _annotate(self) -> tuple[natspec.NatspecOutput, vy_ast.Module]:
@@ -271,7 +284,7 @@ class CompilerData:
     def bytecode(self) -> bytes:
         metadata = None
         if not self.no_bytecode_metadata:
-            metadata = bytes.fromhex(self.resolved_imports.integrity_sum)
+            metadata = bytes.fromhex(self.integrity_sum)
         return generate_bytecode(self.assembly, compiler_metadata=metadata)
 
     @cached_property
diff --git a/vyper/semantics/analysis/imports.py b/vyper/semantics/analysis/imports.py
index 4f8daefa..bcd62feb 100644
--- a/vyper/semantics/analysis/imports.py
+++ b/vyper/semantics/analysis/imports.py
@@ -79,14 +79,14 @@ class ImportAnalyzer:
 
         self.seen: set[int] = set()
 
-        self.integrity_sum = None
+        self._integrity_sum = None
 
         # should be all system paths + topmost module path
         self.absolute_search_paths = input_bundle.search_paths.copy()
 
     def resolve_imports(self, module_ast: vy_ast.Module):
         self._resolve_imports_r(module_ast)
-        self.integrity_sum = self._calculate_integrity_sum_r(module_ast)
+        self._integrity_sum = self._calculate_integrity_sum_r(module_ast)
 
     def _calculate_integrity_sum_r(self, module_ast: vy_ast.Module):
         acc = [sha256sum(module_ast.full_source_code)]


Test patch:
diff --git a/tests/unit/cli/storage_layout/test_storage_layout_overrides.py b/tests/unit/cli/storage_layout/test_storage_layout_overrides.py
index f02a8471..d2c495a2 100644
--- a/tests/unit/cli/storage_layout/test_storage_layout_overrides.py
+++ b/tests/unit/cli/storage_layout/test_storage_layout_overrides.py
@@ -2,6 +2,7 @@ import re
 
 import pytest
 
+from vyper.cli.vyper_json import compile_json
 from vyper.compiler import compile_code
 from vyper.evm.opcodes import version_check
 from vyper.exceptions import StorageLayoutException
@@ -13,7 +14,7 @@ a: uint256
 b: uint256"""
 
     storage_layout_overrides = {
-        "a": {"type": "uint256", "slot": 1, "n_slots": 1},
+        "a": {"type": "uint256", "slot": 5, "n_slots": 1},
         "b": {"type": "uint256", "slot": 0, "n_slots": 1},
     }
 
@@ -26,6 +27,31 @@ b: uint256"""
     assert out["layout"] == expected_output
 
 
+def test_storage_layout_overrides_json():
+    code = """
+a: uint256
+b: uint256"""
+
+    storage_layout_overrides = {
+        "a": {"type": "uint256", "slot": 5, "n_slots": 1},
+        "b": {"type": "uint256", "slot": 0, "n_slots": 1},
+    }
+
+    input_json = {
+        "language": "Vyper",
+        "sources": {"contracts/foo.vy": {"content": code}},
+        "storage_layout_overrides": {"contracts/foo.vy": storage_layout_overrides},
+        "settings": {"outputSelection": {"*": ["*"]}},
+    }
+
+    out = compile_code(
+        code, output_formats=["layout"], storage_layout_override=storage_layout_overrides
+    )
+    assert (
+        compile_json(input_json)["contracts"]["contracts/foo.vy"]["foo"]["layout"] == out["layout"]
+    )
+
+
 def test_storage_layout_for_more_complex():
     code = """
 foo: HashMap[address, uint256]
diff --git a/tests/unit/cli/vyper_compile/test_compile_files.py b/tests/unit/cli/vyper_compile/test_compile_files.py
index 0fd938d5..a1f5ca09 100644
--- a/tests/unit/cli/vyper_compile/test_compile_files.py
+++ b/tests/unit/cli/vyper_compile/test_compile_files.py
@@ -1,18 +1,27 @@
 import contextlib
+import json
 import sys
+import warnings
 import zipfile
 from pathlib import Path
 
 import pytest
 
+from vyper.cli.compile_archive import compiler_data_from_zip
 from vyper.cli.vyper_compile import compile_files
-from vyper.cli.vyper_json import compile_json
+from vyper.cli.vyper_json import compile_from_input_dict, compile_json
 from vyper.compiler import INTERFACE_OUTPUT_FORMATS, OUTPUT_FORMATS
 from vyper.compiler.input_bundle import FilesystemInputBundle
 from vyper.compiler.output_bundle import OutputBundle
 from vyper.compiler.phases import CompilerData
 from vyper.utils import sha256sum
 
+TAMPERED_INTEGRITY_SUM = sha256sum("tampered integrity sum")
+
+INTEGRITY_WARNING = f"Mismatched integrity sum! Expected {TAMPERED_INTEGRITY_SUM}"
+INTEGRITY_WARNING += " but got {integrity}."  # noqa: FS003
+INTEGRITY_WARNING += " (This likely indicates a corrupted archive)"
+
 
 def test_combined_json_keys(chdir_tmp_path, make_file):
     make_file("bar.vy", "")
@@ -298,6 +307,9 @@ def foo() -> uint256:
 import lib
 import jsonabi
 
+a: uint256
+b: uint256
+
 @external
 def foo() -> uint256:
     return lib.foo()
@@ -306,28 +318,52 @@ def foo() -> uint256:
 def bar(x: uint256) -> uint256:
     return extcall jsonabi(msg.sender).test_json(x)
     """
+    storage_layout_overrides = {
+        "a": {"type": "uint256", "n_slots": 1, "slot": 5},
+        "b": {"type": "uint256", "n_slots": 1, "slot": 0},
+    }
+    storage_layout_source = json.dumps(storage_layout_overrides)
+
     tmpdir = tmp_path_factory.mktemp("fake-package")
     with open(tmpdir / "lib.vy", "w") as f:
         f.write(library_source)
     with open(tmpdir / "jsonabi.json", "w") as f:
         f.write(json_source)
+    with open(tmpdir / "layout.json", "w") as f:
+        f.write(storage_layout_source)
 
     contract_file = make_file("contract.vy", contract_source)
 
-    return (tmpdir, tmpdir / "lib.vy", tmpdir / "jsonabi.json", contract_file)
+    contract_hash = sha256sum(contract_source)
+    library_hash = sha256sum(library_source)
+    jsonabi_hash = sha256sum(json_source)
+    resolved_imports_hash = sha256sum(contract_hash + sha256sum(library_hash) + jsonabi_hash)
+    storage_layout_hash = sha256sum(storage_layout_source)
+    expected_integrity = sha256sum(storage_layout_hash + resolved_imports_hash)
+
+    return (
+        tmpdir,
+        tmpdir / "lib.vy",
+        tmpdir / "jsonabi.json",
+        tmpdir / "layout.json",
+        contract_file,
+        expected_integrity,
+    )
 
 
 def test_import_sys_path(input_files):
-    tmpdir, _, _, contract_file = input_files
+    tmpdir, _, _, _, contract_file, _ = input_files
     with mock_sys_path(tmpdir):
         assert compile_files([contract_file], ["combined_json"]) is not None
 
 
 def test_archive_output(input_files):
-    tmpdir, _, _, contract_file = input_files
+    tmpdir, library_file, jsonabi_file, storage_layout_path, contract_file, integrity = input_files
     search_paths = [".", tmpdir]
 
-    s = compile_files([contract_file], ["archive"], paths=search_paths)
+    s = compile_files(
+        [contract_file], ["archive"], paths=search_paths, storage_layout_paths=[storage_layout_path]
+    )
     archive_bytes = s[contract_file]["archive"]
 
     archive_path = Path("foo.zip")
@@ -337,13 +373,28 @@ def test_archive_output(input_files):
     assert zipfile.is_zipfile(archive_path)
 
     # compare compiling the two input bundles
-    out = compile_files([contract_file], ["integrity", "bytecode"], paths=search_paths)
-    out2 = compile_files([archive_path], ["integrity", "bytecode"])
+    out = compile_files(
+        [contract_file],
+        ["integrity", "bytecode", "layout"],
+        paths=search_paths,
+        storage_layout_paths=[storage_layout_path],
+    )
+    out2 = compile_files([archive_path], ["integrity", "bytecode", "layout"])
     assert out[contract_file] == out2[archive_path]
 
+    # tamper with the integrity sum
+    archive_compiler_data = compiler_data_from_zip(archive_path, None, False)
+    archive_compiler_data.expected_integrity_sum = TAMPERED_INTEGRITY_SUM
+
+    with warnings.catch_warnings(record=True) as w:
+        assert archive_compiler_data.integrity_sum is not None
+
+    assert len(w) == 1, [s.message for s in w]
+    assert str(w[0].message).startswith(INTEGRITY_WARNING.format(integrity=integrity))
+
 
 def test_archive_b64_output(input_files):
-    tmpdir, _, _, contract_file = input_files
+    tmpdir, _, _, _, contract_file, _ = input_files
     search_paths = [".", tmpdir]
 
     out = compile_files(
@@ -362,7 +413,7 @@ def test_archive_b64_output(input_files):
 
 
 def test_archive_compile_options(input_files):
-    tmpdir, _, _, contract_file = input_files
+    tmpdir, _, _, _, contract_file, _ = input_files
     search_paths = [".", tmpdir]
 
     options = ["abi_python", "json", "ast", "annotated_ast", "ir_json"]
@@ -417,7 +468,7 @@ format_options = [
 
 
 def test_compile_vyz_with_options(input_files):
-    tmpdir, _, _, contract_file = input_files
+    tmpdir, _, _, _, contract_file, _ = input_files
     search_paths = [".", tmpdir]
 
     for option in format_options:
@@ -446,7 +497,7 @@ def test_compile_vyz_with_options(input_files):
 
 
 def test_archive_compile_simultaneous_options(input_files):
-    tmpdir, _, _, contract_file = input_files
+    tmpdir, _, _, _, contract_file, _ = input_files
     search_paths = [".", tmpdir]
 
     for option in format_options:
@@ -461,40 +512,55 @@ def test_archive_compile_simultaneous_options(input_files):
 
 
 def test_solc_json_output(input_files):
-    tmpdir, _, _, contract_file = input_files
+    tmpdir, _, _, storage_layout_path, contract_file, integrity = input_files
     search_paths = [".", tmpdir]
 
-    out = compile_files([contract_file], ["solc_json"], paths=search_paths)
-
+    out = compile_files(
+        [contract_file],
+        ["solc_json"],
+        paths=search_paths,
+        storage_layout_paths=[storage_layout_path],
+    )
     json_input = out[contract_file]["solc_json"]
 
     # check that round-tripping solc_json thru standard json produces
     # the same as compiling directly
     json_out = compile_json(json_input)["contracts"]["contract.vy"]
     json_out_bytecode = json_out["contract"]["evm"]["bytecode"]["object"]
+    json_out_layout = json_out["contract"]["layout"]["storage_layout"]
 
-    out2 = compile_files([contract_file], ["integrity", "bytecode"], paths=search_paths)
+    out2 = compile_files(
+        [contract_file],
+        ["integrity", "bytecode", "layout"],
+        paths=search_paths,
+        storage_layout_paths=[storage_layout_path],
+    )
 
     assert out2[contract_file]["bytecode"] == json_out_bytecode
+    assert out2[contract_file]["layout"]["storage_layout"] == json_out_layout
+
+    # tamper with the integrity sum
+    json_input["integrity"] = TAMPERED_INTEGRITY_SUM
+    _, warn_data = compile_from_input_dict(json_input)
+
+    w = warn_data[Path("contract.vy")]
+    assert len(w) == 1, [s.message for s in w]
+    assert str(w[0].message).startswith(INTEGRITY_WARNING.format(integrity=integrity))
 
 
 # maybe this belongs in tests/unit/compiler?
 def test_integrity_sum(input_files):
-    tmpdir, library_file, jsonabi_file, contract_file = input_files
+    tmpdir, library_file, jsonabi_file, storage_layout_path, contract_file, integrity = input_files
     search_paths = [".", tmpdir]
 
-    out = compile_files([contract_file], ["integrity"], paths=search_paths)
-
-    with library_file.open() as f, contract_file.open() as g, jsonabi_file.open() as h:
-        library_contents = f.read()
-        contract_contents = g.read()
-        jsonabi_contents = h.read()
+    out = compile_files(
+        [contract_file],
+        ["integrity"],
+        paths=search_paths,
+        storage_layout_paths=[storage_layout_path],
+    )
 
-    contract_hash = sha256sum(contract_contents)
-    library_hash = sha256sum(library_contents)
-    jsonabi_hash = sha256sum(jsonabi_contents)
-    expected = sha256sum(contract_hash + sha256sum(library_hash) + jsonabi_hash)
-    assert out[contract_file]["integrity"] == expected
+    assert out[contract_file]["integrity"] == integrity
 
 
 # does this belong in tests/unit/compiler?
diff --git a/tests/unit/cli/vyper_json/test_compile_json.py b/tests/unit/cli/vyper_json/test_compile_json.py
index 9044148a..f921d250 100644
--- a/tests/unit/cli/vyper_json/test_compile_json.py
+++ b/tests/unit/cli/vyper_json/test_compile_json.py
@@ -20,6 +20,9 @@ import contracts.ibar as IBar
 
 import contracts.library as library
 
+a: uint256
+b: uint256
+
 @external
 def foo(a: address) -> bool:
     return extcall IBar(a).bar(1)
@@ -29,16 +32,29 @@ def baz() -> uint256:
     return self.balance + library.foo()
 """
 
+FOO_STORAGE_LAYOUT_OVERRIDES = {
+    "a": {"type": "uint256", "n_slots": 1, "slot": 5},
+    "b": {"type": "uint256", "n_slots": 1, "slot": 0},
+}
+
 BAR_CODE = """
 import contracts.ibar as IBar
 
 implements: IBar
 
+c: uint256
+d: uint256
+
 @external
 def bar(a: uint256) -> bool:
     return True
 """
 
+BAR_STORAGE_LAYOUT_OVERRIDES = {
+    "c": {"type": "uint256", "n_slots": 1, "slot": 13},
+    "d": {"type": "uint256", "n_slots": 1, "slot": 7},
+}
+
 BAR_VYI = """
 @external
 def bar(a: uint256) -> bool:
@@ -89,6 +105,10 @@ def input_json(optimize, evm_version, experimental_codegen, debug):
             "experimentalCodegen": experimental_codegen,
             "debug": debug,
         },
+        "storage_layout_overrides": {
+            "contracts/foo.vy": FOO_STORAGE_LAYOUT_OVERRIDES,
+            "contracts/bar.vy": BAR_STORAGE_LAYOUT_OVERRIDES,
+        },
     }
 
 
@@ -128,7 +148,10 @@ def test_compile_json(input_json, input_bundle):
     del output_formats["cfg"]
     del output_formats["cfg_runtime"]
     foo = compile_from_file_input(
-        foo_input, output_formats=output_formats, input_bundle=input_bundle
+        foo_input,
+        output_formats=output_formats,
+        input_bundle=input_bundle,
+        storage_layout_override=FOO_STORAGE_LAYOUT_OVERRIDES,
     )
 
     library_input = input_bundle.load_file("contracts/library.vy")
@@ -138,7 +161,10 @@ def test_compile_json(input_json, input_bundle):
 
     bar_input = input_bundle.load_file("contracts/bar.vy")
     bar = compile_from_file_input(
-        bar_input, output_formats=output_formats, input_bundle=input_bundle
+        bar_input,
+        output_formats=output_formats,
+        input_bundle=input_bundle,
+        storage_layout_override=BAR_STORAGE_LAYOUT_OVERRIDES,
     )
 
     compile_code_results = {
@@ -171,6 +197,7 @@ def test_compile_json(input_json, input_bundle):
             "interface": data["interface"],
             "ir": data["ir_dict"],
             "userdoc": data["userdoc"],
+            "layout": data["layout"],
             "metadata": data["metadata"],
             "evm": {
                 "bytecode": {
@@ -218,7 +245,16 @@ def test_different_outputs(input_bundle, input_json):
 
     foo = contracts["contracts/foo.vy"]["foo"]
     bar = contracts["contracts/bar.vy"]["bar"]
-    assert sorted(bar.keys()) == ["abi", "devdoc", "evm", "interface", "ir", "metadata", "userdoc"]
+    assert sorted(bar.keys()) == [
+        "abi",
+        "devdoc",
+        "evm",
+        "interface",
+        "ir",
+        "layout",
+        "metadata",
+        "userdoc",
+    ]
 
     assert sorted(foo.keys()) == ["evm"]
 
@@ -270,6 +306,14 @@ def test_exc_handler_to_dict_compiler(input_json):
     assert error["type"] == "TypeMismatch"
 
 
+def test_unknown_storage_layout_overrides(input_json):
+    unknown_contract_path = "contracts/baz.vy"
+    input_json["storage_layout_overrides"] = {unknown_contract_path: FOO_STORAGE_LAYOUT_OVERRIDES}
+    with pytest.raises(JSONError) as e:
+        compile_json(input_json)
+    assert e.value.args[0] == f"unknown target for storage layout override: {unknown_contract_path}"
+
+
 def test_source_ids_increment(input_json):
     input_json["settings"]["outputSelection"] = {"*": ["ast", "evm.deployedBytecode.sourceMap"]}
     result = compile_json(input_json)