fix[codegen]: fix some hardcoded references to STORAGE location

vyper/ast/nodes.pyi

@@ -234,6 +234,7 @@ class VariableDecl(VyperNode):
     is_constant: bool = ...
     is_public: bool = ...
     is_immutable: bool = ...
+    is_transient: bool = ...
     _expanded_getter: FunctionDef = ...
 
 class AugAssign(VyperNode):

vyper/codegen/core.py

@@ -150,10 +150,10 @@ def make_byte_array_copier(dst, src):
         return STORE(dst, 0)
 
     with src.cache_when_complex("src") as (b1, src):
-        has_storage = STORAGE in (src.location, dst.location)
+        no_copy_opcode = any(not loc.has_copy_opcode for loc in (src.location, dst.location))
         is_memory_copy = dst.location == src.location == MEMORY
         batch_uses_identity = is_memory_copy and not version_check(begin="cancun")
-        if src.typ.maxlen <= 32 and (has_storage or batch_uses_identity):
+        if src.typ.maxlen <= 32 and (no_copy_opcode or batch_uses_identity):
             # it's cheaper to run two load/stores instead of copy_bytes
 
             ret = ["seq"]
@@ -934,8 +934,9 @@ def _complex_make_setter(left, right):
         assert left.encoding == Encoding.VYPER
         len_ = left.typ.memory_bytes_required
 
-        has_storage = STORAGE in (left.location, right.location)
-        if has_storage:
+        # locations with no dedicated copy opcode
+        # (i.e. storage and transient storage)
+        if any(not loc.has_copy_opcode for loc in (left.location, right.location)):
             if _opt_codesize():
                 # assuming PUSH2, a single sstore(dst (sload src)) is 8 bytes,
                 # sstore(add (dst ofst), (sload (add (src ofst)))) is 16 bytes,

vyper/codegen/stmt.py

@@ -17,7 +17,7 @@
 )
 from vyper.codegen.expr import Expr
 from vyper.codegen.return_ import make_return_stmt
-from vyper.evm.address_space import MEMORY, STORAGE
+from vyper.evm.address_space import MEMORY
 from vyper.exceptions import CodegenPanic, StructureException, TypeCheckFailure, tag_exceptions
 from vyper.semantics.types import DArrayT
 from vyper.semantics.types.shortcuts import UINT256_T
@@ -318,12 +318,14 @@ def _get_target(self, target):
         if isinstance(target, vy_ast.Tuple):
             target = Expr(target, self.context).ir_node
             for node in target.args:
-                if (node.location == STORAGE and self.context.is_constant()) or not node.mutable:
+                if (
+                    node.location.word_addressable and self.context.is_constant()
+                ) or not node.mutable:
                     raise TypeCheckFailure(f"Failed constancy check\n{_dbg_expr}")
             return target
 
         target = Expr.parse_pointer_expr(target, self.context)
-        if (target.location == STORAGE and self.context.is_constant()) or not target.mutable:
+        if (target.location.word_addressable and self.context.is_constant()) or not target.mutable:
             raise TypeCheckFailure(f"Failed constancy check\n{_dbg_expr}")
         return target
 

vyper/evm/address_space.py

@@ -27,6 +27,7 @@ class AddrSpace:
     load_op: str
     # TODO maybe make positional instead of defaulting to None
     store_op: Optional[str] = None
+    has_copy_opcode: bool = True  # has a dedicated opcode to copy to memory
 
     @property
     def word_addressable(self) -> bool:
@@ -43,8 +44,8 @@ def word_addressable(self) -> bool:
 # MEMORY = Memory()
 
 MEMORY = AddrSpace("memory", 32, "mload", "mstore")
-STORAGE = AddrSpace("storage", 1, "sload", "sstore")
-TRANSIENT = AddrSpace("transient", 1, "tload", "tstore")
+STORAGE = AddrSpace("storage", 1, "sload", "sstore", has_copy_opcode=False)
+TRANSIENT = AddrSpace("transient", 1, "tload", "tstore", has_copy_opcode=False)
 CALLDATA = AddrSpace("calldata", 32, "calldataload")
 # immutables address space: "immutables" section of memory
 # which is read-write in deploy code but then gets turned into

vyper/semantics/analysis/module.py

@@ -51,8 +51,8 @@
     check_modifiability,
     get_exact_type_from_node,
     get_expr_info,
+    location_from_decl,
 )
-from vyper.semantics.data_locations import DataLocation
 from vyper.semantics.namespace import Namespace, get_namespace, override_global_namespace
 from vyper.semantics.types import EventT, FlagT, InterfaceT, StructT
 from vyper.semantics.types.function import ContractFunctionT
@@ -648,15 +648,7 @@ def visit_VariableDecl(self, node):
                 )
                 raise ImmutableViolation(message, node)
 
-        data_loc = (
-            DataLocation.CODE
-            if node.is_immutable
-            else DataLocation.UNSET
-            if node.is_constant
-            else DataLocation.TRANSIENT
-            if node.is_transient
-            else DataLocation.STORAGE
-        )
+        data_loc = location_from_decl(node)
 
         modifiability = (
             Modifiability.RUNTIME_CONSTANT

vyper/semantics/analysis/utils.py

@@ -19,6 +19,7 @@
 from vyper.semantics import types
 from vyper.semantics.analysis.base import ExprInfo, Modifiability, ModuleInfo, VarAccess, VarInfo
 from vyper.semantics.analysis.levenshtein_utils import get_levenshtein_error_suggestions
+from vyper.semantics.data_locations import DataLocation
 from vyper.semantics.namespace import get_namespace
 from vyper.semantics.types.base import TYPE_T, VyperType
 from vyper.semantics.types.bytestrings import BytesT, StringT
@@ -681,3 +682,23 @@ def check_modifiability(node: vy_ast.ExprNode, modifiability: Modifiability) ->
 
     info = get_expr_info(node)
     return info.modifiability <= modifiability
+
+
+def location_from_decl(node: vy_ast.VariableDecl) -> DataLocation:
+    """
+    Extract the data location from a variable declaration node.
+    """
+
+    assert isinstance(node, vy_ast.VariableDecl)
+
+    data_loc = (
+        DataLocation.CODE
+        if node.is_immutable
+        else DataLocation.UNSET
+        if node.is_constant
+        else DataLocation.TRANSIENT
+        if node.is_transient
+        else DataLocation.STORAGE
+    )
+
+    return data_loc

vyper/semantics/types/function.py

@@ -27,6 +27,7 @@
 from vyper.semantics.analysis.utils import (
     check_modifiability,
     get_exact_type_from_node,
+    location_from_decl,
     uses_state,
     validate_expected_type,
 )
@@ -460,7 +461,13 @@ def getter_from_VariableDecl(cls, node: vy_ast.VariableDecl) -> "ContractFunctio
         """
         if not node.is_public:
             raise CompilerPanic("getter generated for non-public function")
-        type_ = type_from_annotation(node.annotation, DataLocation.STORAGE)
+
+        data_loc = location_from_decl(node)
+
+        assert data_loc not in (DataLocation.MEMORY, DataLocation.CALLDATA)
+
+        type_ = type_from_annotation(node.annotation, data_loc)
+
         arguments, return_type = type_.getter_signature
         args = []
         for i, item in enumerate(arguments):

vyper/semantics/types/subscriptable.py

@@ -46,7 +46,7 @@ class HashMapT(_SubscriptableT):
 
     _equality_attrs = ("key_type", "value_type")
 
-    # disallow everything but storage
+    # disallow everything but storage or transient
     _invalid_locations = (
         DataLocation.UNSET,
         DataLocation.CALLDATA,
@@ -84,10 +84,11 @@ def from_annotation(cls, node: vy_ast.Subscript) -> "HashMapT":
             )
 
         k_ast, v_ast = node.slice.elements
-        key_type = type_from_annotation(k_ast, DataLocation.STORAGE)
+        key_type = type_from_annotation(k_ast)
         if not key_type._as_hashmap_key:
             raise InvalidType("can only use primitive types as HashMap key!", k_ast)
 
+        # TODO: thread through actual location - might also be TRANSIENT
         value_type = type_from_annotation(v_ast, DataLocation.STORAGE)
 
         return cls(key_type, value_type)