fix[codegen]: fix transient codegen for slice and extract32

tests/functional/builtins/codegen/test_extract32.py

@@ -1,6 +1,20 @@
-def test_extract32_extraction(tx_failed, get_contract_with_gas_estimation):
-    extract32_code = """
-y: Bytes[100]
+import pytest
+
+from tests.utils import wrap_typ_with_storage_loc
+
+from vyper.evm.opcodes import version_check
+
+
+@pytest.mark.parametrize('location', ["storage", "transient"])
+def test_extract32_extraction(
+    tx_failed,
+    get_contract_with_gas_estimation,
+    location
+):
+    if location == "transient" and not version_check(begin="cancun"):
+        pytest.skip("Skipping test as storage_location is 'transient' and EVM version is pre-Cancun")
+    extract32_code = f"""
+y: {wrap_typ_with_storage_loc("Bytes[100]", location)}
 @external
 def extrakt32(inp: Bytes[100], index: uint256) -> bytes32:
     return extract32(inp, index)
@@ -43,8 +57,6 @@ def extrakt32_storage(index: uint256, inp: Bytes[100]) -> bytes32:
             with tx_failed():
                 c.extrakt32(S, i)
 
-    print("Passed bytes32 extraction test")
-
 
 def test_extract32_code(tx_failed, get_contract_with_gas_estimation):
     extract32_code = """
@@ -84,5 +96,3 @@ def foq(inp: Bytes[32]) -> address:
 
     with tx_failed():
         c.foq(b"crow" * 8)
-
-    print("Passed extract32 test")

tests/functional/builtins/codegen/test_slice.py

@@ -2,9 +2,12 @@
 import pytest
 from hypothesis import given, settings
 
+from tests.utils import wrap_typ_with_storage_loc
+
 from vyper.compiler import compile_code
 from vyper.compiler.settings import OptimizationLevel, Settings
 from vyper.exceptions import ArgumentException, TypeMismatch
+from vyper.evm.opcodes import version_check
 
 _fun_bytes32_bounds = [(0, 32), (3, 29), (27, 5), (0, 5), (5, 3), (30, 2)]
 
@@ -93,7 +96,7 @@ def _get_contract():
         assert c.do_splice() == bytesdata[start : start + length]
 
 
-@pytest.mark.parametrize("location", ("storage", "calldata", "memory", "literal", "code"))
+@pytest.mark.parametrize("location", ["storage", "transient", "calldata", "memory", "literal", "code"])
 @pytest.mark.parametrize("use_literal_start", (True, False))
 @pytest.mark.parametrize("use_literal_length", (True, False))
 @pytest.mark.parametrize("opt_level", list(OptimizationLevel))
@@ -112,10 +115,18 @@ def test_slice_bytes_fuzz(
     use_literal_length,
     length_bound,
 ):
+    if location == "transient" and not version_check(begin="cancun"):
+        pytest.skip("Skipping test as storage_location is 'transient' and EVM version is pre-Cancun")
     preamble = ""
     if location == "memory":
         spliced_code = f"foo: Bytes[{length_bound}] = inp"
         foo = "foo"
+    elif location == "transient" or location == "storage":
+        preamble = f"""
+foo: {wrap_typ_with_storage_loc(f"Bytes[{length_bound}]", location)}
+         """
+        spliced_code = "self.foo = inp"
+        foo = "self.foo"
     elif location == "storage":
         preamble = f"""
 foo: Bytes[{length_bound}]
@@ -194,10 +205,11 @@ def _get_contract():
         assert c.do_slice(bytesdata, start, length) == bytesdata[start:end], code
 
 
-def test_slice_private(get_contract):
+@pytest.mark.parametrize("location", ["storage"] if not version_check(begin="cancun") else ["storage", "transient"])
+def test_slice_private(get_contract, location):
     # test there are no buffer overruns in the slice function
-    code = """
-bytez: public(String[12])
+    code = f"""
+bytez: public({wrap_typ_with_storage_loc("String[12]", location)})
 
 @internal
 def _slice(start: uint256, length: uint256):

tests/utils.py

@@ -19,3 +19,11 @@ def parse_and_fold(source_code):
     ast = vy_ast.parse_to_ast(source_code)
     constant_fold(ast)
     return ast
+
+
+def wrap_typ_with_storage_loc(typ, loc):
+    if loc == "storage":
+        return typ
+    elif loc == "transient":
+        return f"transient({typ})"
+    assert False, f"unreachable storage location {loc}"

vyper/ast/grammar.lark

@@ -15,6 +15,7 @@ module: ( DOCSTRING
         | event_def
         | function_def
         | immutable_def
+        | transient_def
         | exports_decl
         | _NEWLINE )*
 
@@ -47,9 +48,13 @@ constant_def: (constant_private | constant_with_getter) "=" expr
 immutable: "immutable" "(" type ")"
 immutable_def: NAME ":" immutable
 
+// transient definitions
+transient: "transient" "(" type ")"
+transient_def: NAME ":" transient
+
 variable: NAME ":" type
 // NOTE: Temporary until decorators used
-variable_with_getter: NAME ":" "public" "(" (type | immutable) ")"
+variable_with_getter: NAME ":" "public" "(" (type | immutable | transient) ")"
 variable_def: variable | variable_with_getter
 
 // A decorator "wraps" a method, modifying it's context.

vyper/builtins/functions.py

@@ -36,7 +36,7 @@
 from vyper.codegen.expr import Expr
 from vyper.codegen.ir_node import Encoding, scope_multi
 from vyper.codegen.keccak256_helper import keccak256_helper
-from vyper.evm.address_space import MEMORY, STORAGE
+from vyper.evm.address_space import MEMORY
 from vyper.exceptions import (
     ArgumentException,
     CompilerPanic,
@@ -367,7 +367,7 @@ def build_IR(self, expr, args, kwargs, context):
 
             # add 32 bytes to the buffer size bc word access might
             # be unaligned (see below)
-            if src.location == STORAGE:
+            if src.location.word_addressable:
                 buflen += 32
 
             # Get returntype string or bytes
@@ -394,8 +394,8 @@ def build_IR(self, expr, args, kwargs, context):
                 src_data = bytes_data_ptr(src)
 
             # general case. byte-for-byte copy
-            if src.location == STORAGE:
-                # because slice uses byte-addressing but storage
+            if src.location.word_addressable:
+                # because slice uses byte-addressing but storage/tstorage
                 # is word-aligned, this algorithm starts at some number
                 # of bytes before the data section starts, and might copy
                 # an extra word. the pseudocode is:
@@ -828,19 +828,20 @@ class ECMul(_ECArith):
     _precompile = 0x7
 
 
-def _generic_element_getter(op):
+def _generic_element_getter(loc):
     def f(index):
+        scale = loc.word_scale
+        # Indexing is done by words
+        # - add 'scale' to skip the length slot
+        # - for byte-addressable locations multiply by 32
+        # - for word-addressable locations multiply by 1 which will be optimized out
         return IRnode.from_list(
-            [op, ["add", "_sub", ["add", 32, ["mul", 32, index]]]], typ=INT128_T
+            [loc.load_op, ["add", "_sub", ["add", scale, ["mul", scale, index]]]], typ=INT128_T
         )
 
     return f
 
 
-def _storage_element_getter(index):
-    return IRnode.from_list(["sload", ["add", "_sub", ["add", 1, index]]], typ=INT128_T)
-
-
 class Extract32(BuiltinFunctionT):
     _id = "extract32"
     _inputs = [("b", BytesT.any()), ("start", IntegerT.unsigneds())]
@@ -875,18 +876,10 @@ def build_IR(self, expr, args, kwargs, context):
         sub, index = args
         ret_type = kwargs["output_type"]
 
-        # Get length and specific element
-        if sub.location == STORAGE:
-            lengetter = IRnode.from_list(["sload", "_sub"], typ=INT128_T)
-            elementgetter = _storage_element_getter
-
-        else:
-            op = sub.location.load_op
-            lengetter = IRnode.from_list([op, "_sub"], typ=INT128_T)
-            elementgetter = _generic_element_getter(op)
+        lengetter = IRnode.from_list([sub.location.load_op, "_sub"], typ=INT128_T)
+        elementgetter = _generic_element_getter(sub.location)
 
         # TODO rewrite all this with cache_when_complex and bitshifts
-
         # Special case: index known to be a multiple of 32
         if isinstance(index.value, int) and not index.value % 32:
             o = IRnode.from_list(

vyper/evm/address_space.py

@@ -28,6 +28,10 @@ class AddrSpace:
     # TODO maybe make positional instead of defaulting to None
     store_op: Optional[str] = None
 
+    @property
+    def word_addressable(self) -> bool:
+        return self.word_scale == 1
+
 
 # alternative:
 # class Memory(AddrSpace):