Merge pull request #2433 from charles-cooper/export_layout

export the storage layout

tests/cli/outputs/test_storage_layout.py

@@ -0,0 +1,38 @@
+from vyper.compiler import compile_code
+
+
+def test_method_identifiers():
+    code = """
+foo: HashMap[address, uint256]
+baz: Bytes[65]
+bar: uint256
+
+@external
+@nonreentrant("foo")
+def public_foo():
+    pass
+
+@internal
+@nonreentrant("bar")
+def _bar():
+    pass
+
+@external
+@nonreentrant("bar")
+def public_bar():
+    pass
+    """
+
+    out = compile_code(code, output_formats=["layout"],)
+
+    assert out["layout"] == {
+        "nonreentrant.foo": {"type": "nonreentrant lock", "location": "storage", "slot": 0},
+        "nonreentrant.bar": {"type": "nonreentrant lock", "location": "storage", "slot": 2},
+        "foo": {
+            "type": "HashMap[address, uint256][address, uint256]",
+            "location": "storage",
+            "slot": 3,
+        },
+        "baz": {"type": "Bytes[65]", "location": "storage", "slot": 4},
+        "bar": {"type": "uint256", "location": "storage", "slot": 9},
+    }

tests/cli/vyper_compile/test_compile_files.py

@@ -13,6 +13,7 @@ def test_combined_json_keys(tmp_path):
         "bytecode_runtime",
         "abi",
         "source_map",
+        "layout",
         "method_identifiers",
         "userdoc",
         "devdoc",

vyper/builtin_functions/functions.py

@@ -266,7 +266,7 @@ def build_LLL(self, expr, args, kwargs, context):
             # if we are slicing msg.data, the length should
             # be a constant, since msg.data can be of dynamic length
             # we can't use it's length as the maxlen
-            assert isinstance(length.value, int) # sanity check
+            assert isinstance(length.value, int)  # sanity check
             sub_typ_maxlen = length.value
         else:
             sub_typ_maxlen = sub.typ.maxlen

vyper/cli/vyper_compile.py

@@ -30,6 +30,7 @@
 userdoc            - Natspec user documentation
 devdoc             - Natspec developer documentation
 combined_json      - All of the above format options combined as single JSON output
+layout             - Storage layout of a Vyper contract
 ast                - AST in JSON format
 interface          - Vyper interface of a contract
 external_interface - External interface of a contract, used for outside contract calls
@@ -42,6 +43,7 @@
     "bytecode",
     "bytecode_runtime",
     "abi",
+    "layout",
     "source_map",
     "method_identifiers",
     "userdoc",

vyper/cli/vyper_json.py

@@ -29,6 +29,7 @@
     "evm.deployedBytecode.sourceMap": "source_map",
     "interface": "interface",
     "ir": "ir",
+    "layout": "layout",
     "userdoc": "userdoc",
 }
 

vyper/compiler/__init__.py

@@ -15,6 +15,7 @@
 OUTPUT_FORMATS = {
     # requires vyper_module
     "ast_dict": output.build_ast_dict,
+    "layout": output.build_layout_output,
     # requires global_ctx
     "devdoc": output.build_devdoc,
     "userdoc": output.build_userdoc,

vyper/compiler/output.py

@@ -11,6 +11,7 @@
 from vyper.lll import compile_lll
 from vyper.old_codegen.lll_node import LLLnode
 from vyper.semantics.types.function import FunctionVisibility, StateMutability
+from vyper.typing import StorageLayout
 from vyper.warnings import ContractSizeLimitWarning
 
 
@@ -106,6 +107,12 @@ def build_asm_output(compiler_data: CompilerData) -> str:
     return _build_asm(compiler_data.assembly)
 
 
+def build_layout_output(compiler_data: CompilerData) -> StorageLayout:
+    # in the future this might return (non-storage) layout,
+    # for now only storage layout is returned.
+    return compiler_data.storage_layout
+
+
 def _build_asm(asm_list):
     output_string = ""
     skip_newlines = 0

vyper/compiler/phases.py

@@ -9,7 +9,7 @@
 from vyper.old_codegen import parser
 from vyper.old_codegen.global_context import GlobalContext
 from vyper.semantics import set_data_positions, validate_semantics
-from vyper.typing import InterfaceImports
+from vyper.typing import InterfaceImports, StorageLayout
 
 
 class CompilerData:
@@ -88,10 +88,21 @@ def vyper_module(self) -> vy_ast.Module:
     @property
     def vyper_module_folded(self) -> vy_ast.Module:
         if not hasattr(self, "_vyper_module_folded"):
-            self._vyper_module_folded = generate_folded_ast(self.vyper_module, self.interface_codes)
+            self._vyper_module_folded, self._storage_layout = generate_folded_ast(
+                self.vyper_module, self.interface_codes
+            )
 
         return self._vyper_module_folded
 
+    @property
+    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
+            )
+
+        return self._storage_layout
+
     @property
     def global_ctx(self) -> GlobalContext:
         if not hasattr(self, "_global_ctx"):
@@ -165,7 +176,7 @@ 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]
-) -> vy_ast.Module:
+) -> Tuple[vy_ast.Module, StorageLayout]:
     """
     Perform constant folding operations on the Vyper AST.
 
@@ -178,16 +189,18 @@ def generate_folded_ast(
     -------
     vy_ast.Module
         Folded Vyper AST
+    StorageLayout
+        Layout of variables in storage
     """
     vy_ast.validation.validate_literal_nodes(vyper_module)
 
     vyper_module_folded = copy.deepcopy(vyper_module)
     vy_ast.folding.fold(vyper_module_folded)
     validate_semantics(vyper_module_folded, interface_codes)
     vy_ast.expansion.expand_annotated_ast(vyper_module_folded)
-    set_data_positions(vyper_module_folded)
+    symbol_tables = set_data_positions(vyper_module_folded)
 
-    return vyper_module_folded
+    return vyper_module_folded, symbol_tables
 
 
 def generate_global_context(

vyper/evm/opcodes.py

@@ -195,7 +195,7 @@
     "ASSERT": (None, 1, 0, 85),
     "ASSERT_UNREACHABLE": (None, 1, 0, 17),
     "PASS": (None, 0, 0, 0),
-    "DUMMY": (None, 0, 1, 0), # tell LLL that no, there really is a stack item here
+    "DUMMY": (None, 0, 1, 0),  # tell LLL that no, there really is a stack item here
     "BREAK": (None, 0, 0, 20),
     "CONTINUE": (None, 0, 0, 20),
     "SHA3_32": (None, 1, 1, 72),

vyper/ovm/transpile_lll.py

@@ -657,7 +657,9 @@ def generate_label(opcode):
             + [lll for lll in reversed(rewritten_args)]
             + [["goto", subroutine.subroutine_label()]]
             + [["label", label]]
-            + ["dummy"] if subroutine.evm_returns else ["pass"]
+            + ["dummy"]
+            if subroutine.evm_returns
+            else ["pass"]
         )
 
     lll_ret = [lll_node.value] + rewritten_args

vyper/semantics/validation/data_positions.py

@@ -1,10 +1,12 @@
+# TODO this doesn't really belong in "validation"
 import math
 
 from vyper import ast as vy_ast
 from vyper.semantics.types.bases import StorageSlot
+from vyper.typing import StorageLayout
 
 
-def set_data_positions(vyper_module: vy_ast.Module) -> None:
+def set_data_positions(vyper_module: vy_ast.Module) -> StorageLayout:
     """
     Parse the annotated Vyper AST, determine data positions for all variables,
     and annotate the AST nodes with the position data.
@@ -14,21 +16,34 @@ def set_data_positions(vyper_module: vy_ast.Module) -> None:
     vyper_module : vy_ast.Module
         Top-level Vyper AST node that has already been annotated with type data.
     """
-    set_storage_slots(vyper_module)
+    return set_storage_slots(vyper_module)
 
 
-def set_storage_slots(vyper_module: vy_ast.Module) -> None:
+def set_storage_slots(vyper_module: vy_ast.Module) -> StorageLayout:
     """
     Parse module-level Vyper AST to calculate the layout of storage variables.
+    Returns the layout as a dict of variable name -> variable info
     """
     # Allocate storage slots from 0
     # note storage is word-addressable, not byte-addressable
     storage_slot = 0
 
+    ret = {}
+
     for node in vyper_module.get_children(vy_ast.FunctionDef):
         type_ = node._metadata["type"]
         if type_.nonreentrant is not None:
             type_.set_reentrancy_key_position(StorageSlot(storage_slot))
+
+            # TODO this could have better typing but leave it untyped until
+            # we nail down the format better
+            variable_name = f"nonreentrant.{type_.nonreentrant}"
+            ret[variable_name] = {
+                "type": "nonreentrant lock",
+                "location": "storage",
+                "slot": storage_slot,
+            }
+
             # TODO use one byte - or bit - per reentrancy key
             # requires either an extra SLOAD or caching the value of the
             # location in memory at entrance
@@ -37,12 +52,19 @@ def set_storage_slots(vyper_module: vy_ast.Module) -> None:
     for node in vyper_module.get_children(vy_ast.AnnAssign):
         type_ = node.target._metadata["type"]
         type_.set_position(StorageSlot(storage_slot))
+
+        # this could have better typing but leave it untyped until
+        # we understand the use case better
+        ret[node.target.id] = {"type": str(type_), "location": "storage", "slot": storage_slot}
+
         # CMC 2021-07-23 note that HashMaps get assigned a slot here.
         # I'm not sure if it's safe to avoid allocating that slot
         # for HashMaps because downstream code might use the slot
         # ID as a salt.
         storage_slot += math.ceil(type_.size_in_bytes / 32)
 
+    return ret
+
 
 def set_calldata_offsets(fn_node: vy_ast.FunctionDef) -> None:
     pass

vyper/typing.py

@@ -10,6 +10,7 @@
 ContractCodes = Dict[ContractPath, SourceCode]
 OutputFormats = Sequence[str]
 OutputDict = Dict[ContractPath, OutputFormats]
+StorageLayout = Dict
 
 # Interfaces
 InterfaceAsName = str