feat: add support for modules with variables

tests/functional/codegen/types/test_node_types.py

@@ -12,6 +12,7 @@
 )
 
 # TODO: this module should be merged in with other tests/functional/semantics/types/ tests.
+# and moved to tests/unit/!
 
 
 def test_bytearray_node_type():
@@ -51,17 +52,17 @@ def test_canonicalize_type():
 
 
 def test_type_storage_sizes():
-    assert IntegerT(True, 128).storage_size_in_words == 1
-    assert BytesT(12).storage_size_in_words == 2
-    assert BytesT(33).storage_size_in_words == 3
-    assert SArrayT(IntegerT(True, 128), 10).storage_size_in_words == 10
+    assert IntegerT(True, 128).storage_slots_required == 1
+    assert BytesT(12).storage_slots_required == 2
+    assert BytesT(33).storage_slots_required == 3
+    assert SArrayT(IntegerT(True, 128), 10).storage_slots_required == 10
 
     tuple_ = TupleT([IntegerT(True, 128), DecimalT()])
-    assert tuple_.storage_size_in_words == 2
+    assert tuple_.storage_slots_required == 2
 
     struct_ = StructT("Foo", {"a": IntegerT(True, 128), "b": DecimalT()})
-    assert struct_.storage_size_in_words == 2
+    assert struct_.storage_slots_required == 2
 
     # Don't allow unknown types.
     with raises(AttributeError):
-        _ = int.storage_size_in_words
+        _ = int.storage_slots_required

tests/unit/semantics/types/test_size_in_bytes.py

@@ -11,7 +11,7 @@ def test_base_types(build_node, type_str):
     node = build_node(type_str)
     type_definition = type_from_annotation(node)
 
-    assert type_definition.size_in_bytes == 32
+    assert type_definition._size_in_bytes == 32
 
 
 @pytest.mark.parametrize("type_str", BYTESTRING_TYPES)
@@ -20,7 +20,7 @@ def test_array_value_types(build_node, type_str, length, size):
     node = build_node(f"{type_str}[{length}]")
     type_definition = type_from_annotation(node)
 
-    assert type_definition.size_in_bytes == size
+    assert type_definition._size_in_bytes == size
 
 
 @pytest.mark.parametrize("type_str", BASE_TYPES)
@@ -29,7 +29,7 @@ def test_dynamic_array_lengths(build_node, type_str, length):
     node = build_node(f"DynArray[{type_str}, {length}]")
     type_definition = type_from_annotation(node)
 
-    assert type_definition.size_in_bytes == 32 + length * 32
+    assert type_definition._size_in_bytes == 32 + length * 32
 
 
 @pytest.mark.parametrize("type_str", BASE_TYPES)
@@ -38,7 +38,7 @@ def test_base_types_as_arrays(build_node, type_str, length):
     node = build_node(f"{type_str}[{length}]")
     type_definition = type_from_annotation(node)
 
-    assert type_definition.size_in_bytes == length * 32
+    assert type_definition._size_in_bytes == length * 32
 
 
 @pytest.mark.parametrize("type_str", BASE_TYPES)
@@ -49,4 +49,4 @@ def test_base_types_as_multidimensional_arrays(build_node, type_str, first, seco
 
     type_definition = type_from_annotation(node)
 
-    assert type_definition.size_in_bytes == first * second * 32
+    assert type_definition._size_in_bytes == first * second * 32

vyper/codegen/context.py

@@ -3,10 +3,11 @@
 from dataclasses import dataclass
 from typing import Any, Optional
 
-from vyper.codegen.ir_node import Encoding
-from vyper.evm.address_space import MEMORY, AddrSpace
+import vyper.ast as vy_ast
+from vyper.codegen.ir_node import Encoding, IRnode
+from vyper.evm.address_space import IMMUTABLES, MEMORY, STORAGE, AddrSpace
 from vyper.exceptions import CompilerPanic, StateAccessViolation
-from vyper.semantics.types import VyperType
+from vyper.semantics.types import ModuleT, VyperType
 
 
 class Constancy(enum.Enum):
@@ -48,7 +49,7 @@
 class Context:
     def __init__(
         self,
-        module_ctx,
+        compilation_target,
         memory_allocator,
         vars_=None,
         forvars=None,
@@ -59,9 +60,6 @@
         # In-memory variables, in the form (name, memory location, type)
         self.vars = vars_ or {}
 
-        # Global variables, in the form (name, storage location, type)
-        self.globals = module_ctx.variables
-
         # Variables defined in for loops, e.g. for i in range(6): ...
         self.forvars = forvars or {}
 
@@ -75,8 +73,8 @@
         # Whether we are currently parsing a range expression
         self.in_range_expr = False
 
-        # store module context
-        self.module_ctx = module_ctx
+        # the type information for the current compilation target
+        self.compilation_target: ModuleT = compilation_target
 
         # full function type
         self.func_t = func_t
@@ -94,6 +92,24 @@
         # either the constructor, or called from the constructor
         self.is_ctor_context = is_ctor_context
 
+    def self_ptr(self, location):
+        func_module = self.func_t.ast_def._parent
+        assert isinstance(func_module, vy_ast.Module)
+
+        module_t = func_module._metadata["type"]
+        module_is_compilation_target = module_t == self.compilation_target
+
+        if module_is_compilation_target:
+            # return 0 for the special case where compilation target is self
+            return IRnode.from_list(0, typ=module_t, location=location)
+
+        # otherwise, the function compilation context takes a `self_ptr`
+        # argument in the calling convention
+        if location == STORAGE:
+            return IRnode.from_list("self_ptr_storage", typ=module_t, location=location)
+        if location == IMMUTABLES:
+            return IRnode.from_list("self_ptr_code", typ=module_t, location=location)
+
     def is_constant(self):
         return self.constancy is Constancy.Constant or self.in_assertion or self.in_range_expr
 

vyper/codegen/core.py

@@ -7,6 +7,7 @@
 from vyper.evm.address_space import CALLDATA, DATA, IMMUTABLES, MEMORY, STORAGE, TRANSIENT
 from vyper.evm.opcodes import version_check
 from vyper.exceptions import CompilerPanic, StructureException, TypeCheckFailure, TypeMismatch
+from vyper.semantics.data_locations import DataLocation
 from vyper.semantics.types import (
     AddressT,
     BoolT,
@@ -17,6 +18,7 @@
     HashMapT,
     IntegerT,
     InterfaceT,
+    ModuleT,
     StructT,
     TupleT,
     _BytestringT,
@@ -64,6 +66,17 @@ def is_array_like(typ):
     return ret
 
 
+def data_location_to_addr_space(s: DataLocation):
+    if s == DataLocation.STORAGE:
+        return STORAGE
+    if s == DataLocation.MEMORY:
+        return MEMORY
+    if s == DataLocation.CODE:
+        return IMMUTABLES
+
+    raise CompilerPanic("unreachable")  # pragma: nocover
+
+
 def get_type_for_exact_size(n_bytes):
     """Create a type which will take up exactly n_bytes. Used for allocating internal buffers.
 
@@ -442,6 +455,39 @@ def _getelemptr_abi_helper(parent, member_t, ofst, clamp=True):
     )
 
 
+# get a variable out of a module
+def _get_element_ptr_module(parent, key):
+    # note that this implementation is substantially similar to
+    # the StructT pathway through get_element_ptr_tuplelike and
+    # has potential to be refactored.
+    module_t = parent.typ
+    assert isinstance(module_t, ModuleT)
+
+    assert isinstance(key, str)
+    typ = module_t.variables[key].typ
+    attrs = list(module_t.variables.keys())
+    index = attrs.index(key)
+    annotation = key
+
+    ofst = 0  # offset from parent start
+
+    assert parent.location == STORAGE, parent.location
+
+    for i in range(index):
+        ofst += module_t.variables[attrs[i]].typ.storage_slots_required
+
+    # calculated the same way both ways
+    assert ofst == module_t.variables[key].position.position
+
+    return IRnode.from_list(
+        add_ofst(parent, ofst),
+        typ=typ,
+        location=parent.location,
+        encoding=parent.encoding,
+        annotation=annotation,
+    )
+
+
 # TODO simplify this code, especially the ABI decoding
 def _get_element_ptr_tuplelike(parent, key):
     typ = parent.typ
@@ -485,7 +531,7 @@ def _get_element_ptr_tuplelike(parent, key):
 
     if parent.location.word_addressable:
         for i in range(index):
-            ofst += typ.member_types[attrs[i]].storage_size_in_words
+            ofst += typ.member_types[attrs[i]].storage_slots_required
     elif parent.location.byte_addressable:
         for i in range(index):
             ofst += typ.member_types[attrs[i]].memory_bytes_required
@@ -552,7 +598,7 @@ def _get_element_ptr_array(parent, key, array_bounds_check):
         return _getelemptr_abi_helper(parent, subtype, ofst)
 
     if parent.location.word_addressable:
-        element_size = subtype.storage_size_in_words
+        element_size = subtype.storage_slots_required
     elif parent.location.byte_addressable:
         element_size = subtype.memory_bytes_required
     else:
@@ -590,6 +636,9 @@ def get_element_ptr(parent, key, array_bounds_check=True):
         if is_tuple_like(typ):
             ret = _get_element_ptr_tuplelike(parent, key)
 
+        elif isinstance(typ, ModuleT):
+            ret = _get_element_ptr_module(parent, key)
+
         elif isinstance(typ, HashMapT):
             ret = _get_element_ptr_mapping(parent, key)
 

vyper/codegen/expr.py

@@ -36,6 +36,7 @@
     VyperException,
     tag_exceptions,
 )
+from vyper.semantics.analysis.base import VarInfo
 from vyper.semantics.types import (
     AddressT,
     BoolT,
@@ -169,7 +170,8 @@ def parse_NameConstant(self):
     # Variable names
     def parse_Name(self):
         if self.expr.id == "self":
-            return IRnode.from_list(["address"], typ=AddressT())
+            # TODO: have `self` return a module type
+            return IRnode.from_list(["self"], typ=AddressT())
         elif self.expr.id in self.context.vars:
             var = self.context.vars[self.expr.id]
             ret = IRnode.from_list(
@@ -219,7 +221,7 @@ def parse_Attribute(self):
             return IRnode.from_list(value, typ=typ)
 
         # x.balance: balance of address x
-        if self.expr.attr == "balance":
+        elif self.expr.attr == "balance":
             addr = Expr.parse_value_expr(self.expr.value, self.context)
             if addr.typ == AddressT():
                 if (
@@ -231,6 +233,7 @@ def parse_Attribute(self):
                 else:
                     seq = ["balance", addr]
                 return IRnode.from_list(seq, typ=UINT256_T)
+
         # x.codesize: codesize of address x
         elif self.expr.attr == "codesize" or self.expr.attr == "is_contract":
             addr = Expr.parse_value_expr(self.expr.value, self.context)
@@ -242,14 +245,17 @@ def parse_Attribute(self):
                         eval_code = ["extcodesize", addr]
                     output_type = UINT256_T
                 else:
+                    assert self.expr.attr == "is_contract"
                     eval_code = ["gt", ["extcodesize", addr], 0]
                     output_type = BoolT()
                 return IRnode.from_list(eval_code, typ=output_type)
+
         # x.codehash: keccak of address x
         elif self.expr.attr == "codehash":
             addr = Expr.parse_value_expr(self.expr.value, self.context)
             if addr.typ == AddressT():
                 return IRnode.from_list(["extcodehash", addr], typ=BYTES32_T)
+
         # x.code: codecopy/extcodecopy of address x
         elif self.expr.attr == "code":
             addr = Expr.parse_value_expr(self.expr.value, self.context)
@@ -258,24 +264,12 @@ def parse_Attribute(self):
                 if addr.value == "address":  # for `self.code`
                     return IRnode.from_list(["~selfcode"], typ=BytesT(0))
                 return IRnode.from_list(["~extcode", addr], typ=BytesT(0))
-        # self.x: global attribute
-        elif isinstance(self.expr.value, vy_ast.Name) and self.expr.value.id == "self":
-            varinfo = self.context.globals[self.expr.attr]
-            location = TRANSIENT if varinfo.is_transient else STORAGE
-
-            ret = IRnode.from_list(
-                varinfo.position.position,
-                typ=varinfo.typ,
-                location=location,
-                annotation="self." + self.expr.attr,
-            )
-            ret._referenced_variables = {varinfo}
-
-            return ret
 
         # Reserved keywords
         elif (
-            isinstance(self.expr.value, vy_ast.Name) and self.expr.value.id in ENVIRONMENT_VARIABLES
+            # TODO: use type information here
+            isinstance(self.expr.value, vy_ast.Name)
+            and self.expr.value.id in ENVIRONMENT_VARIABLES
         ):
             key = f"{self.expr.value.id}.{self.expr.attr}"
             if key == "msg.sender":
@@ -327,17 +321,34 @@ def parse_Attribute(self):
                         "chain.id is unavailable prior to istanbul ruleset", self.expr
                     )
                 return IRnode.from_list(["chainid"], typ=UINT256_T)
-        # Other variables
-        else:
-            sub = Expr(self.expr.value, self.context).ir_node
-            # contract type
-            if isinstance(sub.typ, InterfaceT):
-                # MyInterface.address
-                assert self.expr.attr == "address"
-                sub.typ = typ
-                return sub
-            if isinstance(sub.typ, StructT) and self.expr.attr in sub.typ.member_types:
-                return get_element_ptr(sub, self.expr.attr)
+
+        # self.x: global storage variable or immutable
+        if (varinfo := self.expr._metadata.get("variable_access")) is not None:
+            assert isinstance(varinfo, VarInfo)
+
+            # TODO: handle immutables
+            location = TRANSIENT if varinfo.is_transient else STORAGE
+
+            module_ptr = Expr(self.expr.value, self.context).ir_node
+            if module_ptr.value == "self":
+                module_ptr = self.context.self_ptr(location)
+
+            ret = get_element_ptr(module_ptr, self.expr.attr)
+            # TODO: take referenced variables info from analysis
+            ret._referenced_variables = {varinfo}
+            return ret
+
+        # if we have gotten here, it's an instance of an interface or struct
+        sub = Expr(self.expr.value, self.context).ir_node
+
+        if isinstance(sub.typ, InterfaceT):
+            # MyInterface.address
+            assert self.expr.attr == "address"
+            sub.typ = typ
+            return sub
+
+        if isinstance(sub.typ, StructT) and self.expr.attr in sub.typ.member_types:
+            return get_element_ptr(sub, self.expr.attr)
 
     def parse_Subscript(self):
         sub = Expr(self.expr.value, self.context).ir_node

vyper/codegen/function_definitions/common.py

@@ -101,7 +101,7 @@ class InternalFuncIR(FuncIR):
 
 # TODO: should split this into external and internal ir generation?
 def generate_ir_for_function(
-    code: vy_ast.FunctionDef, module_ctx: ModuleT, is_ctor_context: bool = False
+    code: vy_ast.FunctionDef, compilation_target: ModuleT, is_ctor_context: bool = False
 ) -> FuncIR:
     """
     Parse a function and produce IR code for the function, includes:
@@ -133,7 +133,7 @@ def generate_ir_for_function(
 
     context = Context(
         vars_=None,
-        module_ctx=module_ctx,
+        compilation_target=compilation_target,
         memory_allocator=memory_allocator,
         constancy=Constancy.Mutable if func_t.is_mutable else Constancy.Constant,
         func_t=func_t,