From 63910a42cdbf9bc958ac4d67da2541747f5aa743 Mon Sep 17 00:00:00 2001
From: Daniel Schiavini <daniel.schiavini@gmail.com>
Date: Mon, 2 Sep 2024 15:51:40 +0200
Subject: [PATCH 01/20] WIP: VVM internal variables

---
 boa/contracts/vvm/vvm_contract.py        | 238 +++++++++++++++++++++--
 boa/interpret.py                         |   8 +-
 tests/unitary/contracts/vvm/mock_3_10.vy |  13 ++
 tests/unitary/contracts/vvm/test_vvm.py  |  29 ++-
 4 files changed, 259 insertions(+), 29 deletions(-)

diff --git a/boa/contracts/vvm/vvm_contract.py b/boa/contracts/vvm/vvm_contract.py
index 2c2c783a..376e3133 100644
--- a/boa/contracts/vvm/vvm_contract.py
+++ b/boa/contracts/vvm/vvm_contract.py
@@ -1,14 +1,20 @@
 import re
 from functools import cached_property
+from pathlib import Path
 
-from boa.contracts.abi.abi_contract import ABIContractFactory, ABIFunction
+import vvm
+from vyper.utils import method_id
+
+from boa.contracts.abi.abi_contract import ABIContract, ABIContractFactory, ABIFunction
 from boa.environment import Env
+from boa.rpc import to_bytes
+from boa.util.abi import Address
 
 # TODO: maybe this doesn't detect release candidates
 VERSION_RE = re.compile(r"\s*#\s*(pragma\s+version|@version)\s+(\d+\.\d+\.\d+)")
 
 
-# TODO: maybe move this up to vvm?
+# TODO: maybe move this up to VVM?
 def _detect_version(source_code: str):
     res = VERSION_RE.findall(source_code)
     if len(res) < 1:
@@ -17,22 +23,40 @@ def _detect_version(source_code: str):
     return res[0][1]
 
 
-class VVMDeployer:
-    def __init__(self, abi, bytecode, filename):
-        self.abi = abi
-        self.bytecode = bytecode
-        self.filename = filename
-
-    @classmethod
-    def from_compiler_output(cls, compiler_output, filename):
-        abi = compiler_output["abi"]
-        bytecode_nibbles = compiler_output["bytecode"]
-        bytecode = bytes.fromhex(bytecode_nibbles.removeprefix("0x"))
-        return cls(abi, bytecode, filename)
+class VVMDeployer(ABIContractFactory):
+    def __init__(
+        self,
+        name: str,
+        compiler_output: dict,
+        source_code: str,
+        vyper_version: str,
+        filename: str,
+    ):
+        super().__init__(name, compiler_output["abi"], filename)
+        self.compiler_output = compiler_output
+        self.source_code = source_code
+        self.vyper_version = vyper_version
 
     @cached_property
-    def factory(self):
-        return ABIContractFactory.from_abi_dict(self.abi)
+    def bytecode(self):
+        return to_bytes(self.compiler_output["bytecode"])
+
+    @property
+    def layout(self):
+        return self.compiler_output["layout"]
+
+    @classmethod
+    def from_compiler_output(
+        cls,
+        compiler_output: dict,
+        source_code: str,
+        vyper_version: str,
+        filename: str,
+        name: str | None = None,
+    ):
+        if name is None:
+            name = Path(filename).stem
+        return cls(name, compiler_output, source_code, vyper_version, filename)
 
     @cached_property
     def constructor(self):
@@ -58,5 +82,183 @@ def deploy(self, *args, env=None):
     def __call__(self, *args, **kwargs):
         return self.deploy(*args, **kwargs)
 
-    def at(self, address):
-        return self.factory.at(address)
+    def at(self, address: Address | str) -> "VVMContract":
+        """
+        Create an ABI contract object for a deployed contract at `address`.
+        """
+        address = Address(address)
+        contract = VVMContract(
+            self.compiler_output,
+            self.source_code,
+            self.vyper_version,
+            self._name,
+            self._abi,
+            self.functions,
+            address,
+            self.filename,
+        )
+        contract.env.register_contract(address, contract)
+        return contract
+
+
+class VVMContract(ABIContract):
+    def __init__(self, compiler_output, source_code, vyper_version, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.compiler_output = compiler_output
+        self.source_code = source_code
+        self.vyper_version = vyper_version
+
+    @cached_property
+    def bytecode(self):
+        return to_bytes(self.compiler_output["bytecode"])
+
+    @cached_property
+    def bytecode_runtime(self):
+        return to_bytes(self.compiler_output["bytecode_runtime"])
+
+    # def eval(self, code):
+    #     return VVMEval(code, self)()
+
+    @cached_property
+    def _storage(self):
+        def storage():
+            return None
+
+        for name, spec in self.compiler_output["layout"]["storage_layout"].items():
+            setattr(storage, name, VVMStorageVariable(name, spec, self))
+        return storage
+
+    @cached_property
+    def internal(self):
+        def internal():
+            return None
+
+        source = self.source_code.replace("@internal", "@external")
+        result = vvm.compile_source(source, vyper_version=self.vyper_version)
+        for abi_item in result["<stdin>"]["abi"]:
+            if abi_item["type"] == "function" and not isinstance(
+                getattr(self, abi_item["name"], None), ABIFunction
+            ):
+                function = VVMInternalFunction(abi_item, self)
+                setattr(internal, function.name, function)
+        return internal
+
+
+class _VVMInternal(ABIFunction):
+    """
+    An ABI function that temporarily changes the bytecode at the contract's address.
+    """
+
+    @cached_property
+    def _override_bytecode(self) -> bytes:
+        assert isinstance(self.contract, VVMContract)
+        source = "\n".join((self.contract.source_code, self.source_code))
+        compiled = vvm.compile_source(source, vyper_version=self.contract.vyper_version)
+        return to_bytes(compiled["<stdin>"]["bytecode_runtime"])
+
+    @property
+    def source_code(self):
+        raise NotImplementedError  # to be implemented in subclasses
+
+    def __call__(self, *args, **kwargs):
+        env = self.contract.env
+        assert isinstance(self.contract, VVMContract)
+        balance_before = env.get_balance(env.eoa)
+        env.set_code(self.contract.address, self._override_bytecode)
+        env.set_balance(env.eoa, 10**20)
+        try:
+            return super().__call__(*args, **kwargs)
+        finally:
+            env.set_balance(env.eoa, balance_before)
+            env.set_code(self.contract.address, self.contract.bytecode_runtime)
+
+
+class VVMInternalFunction(_VVMInternal):
+    def __init__(self, abi: dict, contract: VVMContract):
+        super().__init__(abi, contract.contract_name)
+        self.contract = contract
+
+    @cached_property
+    def method_id(self) -> bytes:
+        return method_id(f"__boa_internal_{self.name}__" + self.signature)
+
+    @cached_property
+    def source_code(self):
+        fn_args = ", ".join([arg["name"] for arg in self._abi["inputs"]])
+
+        return_sig = ""
+        fn_call = ""
+        if self.return_type:
+            return_sig = f" -> {self.return_type}"
+            fn_call = "return "
+
+        fn_call += f"self.{self.name}({fn_args})"
+        fn_sig = ", ".join(
+            f"{arg['name']}: {arg['type']}" for arg in self._abi["inputs"]
+        )
+        return f"""
+@external
+@payable
+def __boa_internal_{self.name}__({fn_sig}){return_sig}:
+    {fn_call}
+"""
+
+
+class VVMStorageVariable(_VVMInternal):
+    def __init__(self, name, spec, contract):
+        abi = {
+            "anonymous": False,
+            "inputs": [],
+            "outputs": [{"name": name, "type": spec["type"]}],
+            "name": name,
+            "type": "function",
+        }
+
+        if spec["type"].startswith("HashMap"):
+            key_type, value_type = spec["type"][8:-1].split(",")
+            abi["inputs"] = [{"name": "key", "type": key_type}]
+            abi["outputs"] = [{"name": "value", "type": value_type.strip()}]
+
+        super().__init__(abi, contract.contract_name)
+        self.contract = contract
+
+    def get(self, *args):
+        return self.__call__(*args)
+
+    @cached_property
+    def method_id(self) -> bytes:
+        return method_id(f"__boa_private_{self.name}__" + self.signature)
+
+    @cached_property
+    def source_code(self):
+        getter_call = "".join(f"[{i['name']}]" for i in self._abi["inputs"])
+        args_signature = ", ".join(
+            f"{i['name']}: {i['type']}" for i in self._abi["inputs"]
+        )
+        return f"""
+@external
+@payable
+def __boa_private_{self.name}__({args_signature}) -> {self.return_type[0]}:
+    return self.{self.name}{getter_call}
+"""
+
+
+# class VVMEval(_VVMInternal):
+#     def __init__(self, code: str, contract: VVMContract):
+#         typ = detect_expr_type(code, contract)
+#         abi = {
+#             "anonymous": False,
+#             "inputs": [],
+#             "outputs": (
+#                 [{"name": "eval", "type": typ.abi_type.selector_name()}] if typ else []
+#             ),
+#             "name": "__boa_debug__",
+#             "type": "function",
+#         }
+#         super().__init__(abi, contract._name)
+#         self.contract = contract
+#         self.code = code
+#
+#     @cached_property
+#     def source_code(self):
+#         return generate_source_for_arbitrary_stmt(self.code, self.contract)
diff --git a/boa/interpret.py b/boa/interpret.py
index 8386d1fb..ee9e69a9 100644
--- a/boa/interpret.py
+++ b/boa/interpret.py
@@ -216,7 +216,7 @@ def loads_partial(
     version = _detect_version(source_code)
     if version is not None and version != vyper.__version__:
         filename = str(filename)  # help mypy
-        return _loads_partial_vvm(source_code, version, filename)
+        return _loads_partial_vvm(source_code, version, name, filename)
 
     compiler_args = compiler_args or {}
 
@@ -232,14 +232,16 @@ def load_partial(filename: str, compiler_args=None):
         )
 
 
-def _loads_partial_vvm(source_code: str, version: str, filename: str):
+def _loads_partial_vvm(source_code: str, version: str, filename: str, name: str | None):
     # will install the request version if not already installed
     vvm.install_vyper(version=version)
     # TODO: implement caching
     compiled_src = vvm.compile_source(source_code, vyper_version=version)
 
     compiler_output = compiled_src["<stdin>"]
-    return VVMDeployer.from_compiler_output(compiler_output, filename=filename)
+    return VVMDeployer.from_compiler_output(
+        compiler_output, source_code, version, filename, name
+    )
 
 
 def from_etherscan(
diff --git a/tests/unitary/contracts/vvm/mock_3_10.vy b/tests/unitary/contracts/vvm/mock_3_10.vy
index bd23c5a7..339f1ce9 100644
--- a/tests/unitary/contracts/vvm/mock_3_10.vy
+++ b/tests/unitary/contracts/vvm/mock_3_10.vy
@@ -3,7 +3,20 @@
 foo: public(uint256)
 bar: public(uint256)
 
+map: HashMap[address, uint256]
+is_empty: bool
+
 @external
 def __init__(bar: uint256):
     self.foo = 42
     self.bar = bar
+    self.is_empty = True
+
+@external
+def set_map(x: uint256):
+    self._set_map(msg.sender, x)
+
+@internal
+def _set_map(addr: address, x: uint256):
+    self.map[addr] = x
+    self.is_empty = False
diff --git a/tests/unitary/contracts/vvm/test_vvm.py b/tests/unitary/contracts/vvm/test_vvm.py
index 6a2da16f..5c91b7d4 100644
--- a/tests/unitary/contracts/vvm/test_vvm.py
+++ b/tests/unitary/contracts/vvm/test_vvm.py
@@ -1,6 +1,8 @@
 import boa
 
 mock_3_10_path = "tests/unitary/contracts/vvm/mock_3_10.vy"
+with open(mock_3_10_path) as f:
+    mock_3_10_code = f.read()
 
 
 def test_load_partial_vvm():
@@ -12,10 +14,7 @@ def test_load_partial_vvm():
 
 
 def test_loads_partial_vvm():
-    with open(mock_3_10_path) as f:
-        code = f.read()
-
-    contract_deployer = boa.loads_partial(code)
+    contract_deployer = boa.loads_partial(mock_3_10_code)
     contract = contract_deployer.deploy(43)
 
     assert contract.foo() == 42
@@ -30,10 +29,24 @@ def test_load_vvm():
 
 
 def test_loads_vvm():
-    with open(mock_3_10_path) as f:
-        code = f.read()
-
-    contract = boa.loads(code, 43)
+    contract = boa.loads(mock_3_10_code, 43)
 
     assert contract.foo() == 42
     assert contract.bar() == 43
+
+
+def test_vvm_storage():
+    contract = boa.loads(mock_3_10_code, 43)
+    assert contract._storage.is_empty.get()
+    assert contract._storage.map.get(boa.env.eoa) == 0
+    contract.set_map(69)
+    assert not contract._storage.is_empty.get()
+    assert contract._storage.map.get(boa.env.eoa) == 69
+
+
+def test_vvm_internal():
+    contract = boa.loads(mock_3_10_code, 43)
+    assert not hasattr(contract.internal, "set_map")
+    address = boa.env.generate_address()
+    assert contract.internal._set_map(address, 69)
+    assert contract._storage.map.get(address) == 69

From 99bf8602336294c2e5fd7f4af6070955d0bbc18e Mon Sep 17 00:00:00 2001
From: Daniel Schiavini <daniel.schiavini@gmail.com>
Date: Mon, 2 Sep 2024 16:36:48 +0200
Subject: [PATCH 02/20] VVM internal functions via vyper wrapper

---
 boa/contracts/abi/abi_contract.py       |  5 +-
 boa/contracts/base_evm_contract.py      |  3 +-
 boa/contracts/vvm/vvm_contract.py       | 71 ++++++++++++++++++-------
 boa/interpret.py                        | 12 +++--
 tests/unitary/contracts/vvm/test_vvm.py |  2 +-
 5 files changed, 66 insertions(+), 27 deletions(-)

diff --git a/boa/contracts/abi/abi_contract.py b/boa/contracts/abi/abi_contract.py
index 70104c7b..1a357fc4 100644
--- a/boa/contracts/abi/abi_contract.py
+++ b/boa/contracts/abi/abi_contract.py
@@ -1,5 +1,6 @@
 from collections import defaultdict
 from functools import cached_property
+from pathlib import Path
 from typing import Any, Optional, Union
 from warnings import warn
 
@@ -235,7 +236,7 @@ def __init__(
         abi: list[dict],
         functions: list[ABIFunction],
         address: Address,
-        filename: Optional[str] = None,
+        filename: str | Path | None = None,
         env=None,
     ):
         super().__init__(name, env, filename=filename, address=address)
@@ -341,7 +342,7 @@ class ABIContractFactory:
     do any contract deployment.
     """
 
-    def __init__(self, name: str, abi: list[dict], filename: Optional[str] = None):
+    def __init__(self, name: str, abi: list[dict], filename: str | Path | None = None):
         self._name = name
         self._abi = abi
         self.filename = filename
diff --git a/boa/contracts/base_evm_contract.py b/boa/contracts/base_evm_contract.py
index 5e2a45bf..f239a488 100644
--- a/boa/contracts/base_evm_contract.py
+++ b/boa/contracts/base_evm_contract.py
@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from pathlib import Path
 from typing import TYPE_CHECKING, Optional
 
 from eth.abc import ComputationAPI
@@ -25,7 +26,7 @@ def __init__(
         self,
         name: str,
         env: Optional[Env] = None,
-        filename: Optional[str] = None,
+        filename: str | Path | None = None,
         address: Optional[Address] = None,
     ):
         self.env = env or Env.get_singleton()
diff --git a/boa/contracts/vvm/vvm_contract.py b/boa/contracts/vvm/vvm_contract.py
index 376e3133..c08166ff 100644
--- a/boa/contracts/vvm/vvm_contract.py
+++ b/boa/contracts/vvm/vvm_contract.py
@@ -1,8 +1,12 @@
+import json
 import re
 from functools import cached_property
 from pathlib import Path
+from tempfile import NamedTemporaryFile
 
 import vvm
+from vvm.install import get_executable
+from vvm.wrapper import vyper_wrapper
 from vyper.utils import method_id
 
 from boa.contracts.abi.abi_contract import ABIContract, ABIContractFactory, ABIFunction
@@ -30,7 +34,7 @@ def __init__(
         compiler_output: dict,
         source_code: str,
         vyper_version: str,
-        filename: str,
+        filename: str | Path | None = None,
     ):
         super().__init__(name, compiler_output["abi"], filename)
         self.compiler_output = compiler_output
@@ -51,11 +55,11 @@ def from_compiler_output(
         compiler_output: dict,
         source_code: str,
         vyper_version: str,
-        filename: str,
+        filename: str | Path | None = None,
         name: str | None = None,
     ):
         if name is None:
-            name = Path(filename).stem
+            name = Path(filename).stem if filename is not None else "<VVMContract>"
         return cls(name, compiler_output, source_code, vyper_version, filename)
 
     @cached_property
@@ -88,14 +92,14 @@ def at(self, address: Address | str) -> "VVMContract":
         """
         address = Address(address)
         contract = VVMContract(
-            self.compiler_output,
-            self.source_code,
-            self.vyper_version,
-            self._name,
-            self._abi,
-            self.functions,
-            address,
-            self.filename,
+            compiler_output=self.compiler_output,
+            source_code=self.source_code,
+            vyper_version=self.vyper_version,
+            name=self._name,
+            abi=self._abi,
+            functions=self.functions,
+            address=address,
+            filename=self.filename,
         )
         contract.env.register_contract(address, contract)
         return contract
@@ -133,16 +137,30 @@ def internal(self):
         def internal():
             return None
 
-        source = self.source_code.replace("@internal", "@external")
-        result = vvm.compile_source(source, vyper_version=self.vyper_version)
-        for abi_item in result["<stdin>"]["abi"]:
-            if abi_item["type"] == "function" and not isinstance(
-                getattr(self, abi_item["name"], None), ABIFunction
-            ):
-                function = VVMInternalFunction(abi_item, self)
+        result = self._compile_metadata_fn_info()
+        for fn_name, meta in result.items():
+            if meta["visibility"] == "internal":
+                function = VVMInternalFunction(meta, self)
                 setattr(internal, function.name, function)
         return internal
 
+    def _compile_metadata_fn_info(self):
+        # todo: move this to vvm
+        if self.filename is not None:
+            return self._call_vyper(self.filename)
+        with NamedTemporaryFile(suffix=".vy") as f:
+            f.write(self.source_code.encode())
+            f.flush()
+            return self._call_vyper(f.name)
+
+    def _call_vyper(self, filename):
+        stdoutdata, stderrdata, command, proc = vyper_wrapper(
+            vyper_binary=get_executable(self.vyper_version),
+            f="metadata",
+            source_files=[filename],
+        )
+        return json.loads(stdoutdata)["function_info"]
+
 
 class _VVMInternal(ABIFunction):
     """
@@ -174,7 +192,22 @@ def __call__(self, *args, **kwargs):
 
 
 class VVMInternalFunction(_VVMInternal):
-    def __init__(self, abi: dict, contract: VVMContract):
+    def __init__(self, meta: dict, contract: VVMContract):
+        abi = {
+            "anonymous": False,
+            "inputs": [
+                {"name": arg_name, "type": arg_type}
+                for arg_name, arg_type in meta["positional_args"].items()
+            ],
+            "outputs": (
+                [{"name": meta["name"], "type": meta["return_type"]}]
+                if meta["return_type"] != "None"
+                else []
+            ),
+            "stateMutability": meta["mutability"],
+            "name": meta["name"],
+            "type": "function",
+        }
         super().__init__(abi, contract.contract_name)
         self.contract = contract
 
diff --git a/boa/interpret.py b/boa/interpret.py
index ee9e69a9..7db41148 100644
--- a/boa/interpret.py
+++ b/boa/interpret.py
@@ -209,18 +209,17 @@ def loads_partial(
     compiler_args: dict = None,
 ) -> VyperDeployer:
     name = name or "VyperContract"  # TODO handle this upstream in CompilerData
-    filename = filename or "<unknown>"
     if dedent:
         source_code = textwrap.dedent(source_code)
 
     version = _detect_version(source_code)
     if version is not None and version != vyper.__version__:
-        filename = str(filename)  # help mypy
-        return _loads_partial_vvm(source_code, version, name, filename)
+        return _loads_partial_vvm(source_code, version, filename, name)
 
     compiler_args = compiler_args or {}
 
     deployer_class = _get_default_deployer_class()
+    filename = filename or "<unknown>"
     data = compiler_data(source_code, name, filename, deployer_class, **compiler_args)
     return deployer_class(data, filename=filename)
 
@@ -232,7 +231,12 @@ def load_partial(filename: str, compiler_args=None):
         )
 
 
-def _loads_partial_vvm(source_code: str, version: str, filename: str, name: str | None):
+def _loads_partial_vvm(
+    source_code: str,
+    version: str,
+    filename: str | Path | None = None,
+    name: str | None = None,
+):
     # will install the request version if not already installed
     vvm.install_vyper(version=version)
     # TODO: implement caching
diff --git a/tests/unitary/contracts/vvm/test_vvm.py b/tests/unitary/contracts/vvm/test_vvm.py
index 5c91b7d4..72034825 100644
--- a/tests/unitary/contracts/vvm/test_vvm.py
+++ b/tests/unitary/contracts/vvm/test_vvm.py
@@ -48,5 +48,5 @@ def test_vvm_internal():
     contract = boa.loads(mock_3_10_code, 43)
     assert not hasattr(contract.internal, "set_map")
     address = boa.env.generate_address()
-    assert contract.internal._set_map(address, 69)
+    contract.internal._set_map(address, 69)
     assert contract._storage.map.get(address) == 69

From bae493b50bd28dcafe41d5d08ccb86525d8eb239 Mon Sep 17 00:00:00 2001
From: Daniel Schiavini <daniel.schiavini@gmail.com>
Date: Mon, 2 Sep 2024 16:56:34 +0200
Subject: [PATCH 03/20] VVM eval

---
 boa/contracts/vvm/vvm_contract.py       | 49 ++++++++++++++-----------
 tests/unitary/contracts/vvm/test_vvm.py |  7 ++++
 2 files changed, 35 insertions(+), 21 deletions(-)

diff --git a/boa/contracts/vvm/vvm_contract.py b/boa/contracts/vvm/vvm_contract.py
index c08166ff..37c064df 100644
--- a/boa/contracts/vvm/vvm_contract.py
+++ b/boa/contracts/vvm/vvm_contract.py
@@ -120,8 +120,8 @@ def bytecode(self):
     def bytecode_runtime(self):
         return to_bytes(self.compiler_output["bytecode_runtime"])
 
-    # def eval(self, code):
-    #     return VVMEval(code, self)()
+    def eval(self, code, return_type=None):
+        return VVMEval(code, self, return_type)()
 
     @cached_property
     def _storage(self):
@@ -276,22 +276,29 @@ def __boa_private_{self.name}__({args_signature}) -> {self.return_type[0]}:
 """
 
 
-# class VVMEval(_VVMInternal):
-#     def __init__(self, code: str, contract: VVMContract):
-#         typ = detect_expr_type(code, contract)
-#         abi = {
-#             "anonymous": False,
-#             "inputs": [],
-#             "outputs": (
-#                 [{"name": "eval", "type": typ.abi_type.selector_name()}] if typ else []
-#             ),
-#             "name": "__boa_debug__",
-#             "type": "function",
-#         }
-#         super().__init__(abi, contract._name)
-#         self.contract = contract
-#         self.code = code
-#
-#     @cached_property
-#     def source_code(self):
-#         return generate_source_for_arbitrary_stmt(self.code, self.contract)
+class VVMEval(_VVMInternal):
+    def __init__(self, code: str, contract: VVMContract, return_type: str = None):
+        abi = {
+            "anonymous": False,
+            "inputs": [],
+            "outputs": ([{"name": "eval", "type": return_type}] if return_type else []),
+            "name": "__boa_debug__",
+            "type": "function",
+        }
+        super().__init__(abi, contract.contract_name)
+        self.contract = contract
+        self.code = code
+
+    @cached_property
+    def source_code(self):
+        debug_body = self.code
+        return_sig = ""
+        if self.return_type:
+            return_sig = f"-> ({', '.join(self.return_type)})"
+            debug_body = f"return {self.code}"
+        return f"""
+@external
+@payable
+def __boa_debug__() {return_sig}:
+    {debug_body}
+"""
diff --git a/tests/unitary/contracts/vvm/test_vvm.py b/tests/unitary/contracts/vvm/test_vvm.py
index 72034825..9e29b349 100644
--- a/tests/unitary/contracts/vvm/test_vvm.py
+++ b/tests/unitary/contracts/vvm/test_vvm.py
@@ -50,3 +50,10 @@ def test_vvm_internal():
     address = boa.env.generate_address()
     contract.internal._set_map(address, 69)
     assert contract._storage.map.get(address) == 69
+
+
+def test_vvm_eval():
+    contract = boa.loads(mock_3_10_code, 43)
+    assert contract.eval("self.bar", "uint256") == 43
+    assert contract.eval("self.bar = 44") is None
+    assert contract.bar() == 44

From 2fd09c03831e6db940d63b9564d5666108d3a30f Mon Sep 17 00:00:00 2001
From: Daniel Schiavini <daniel.schiavini@gmail.com>
Date: Mon, 2 Sep 2024 20:42:15 +0200
Subject: [PATCH 04/20] Self-review

---
 boa/contracts/abi/abi_contract.py  |  5 +-
 boa/contracts/base_evm_contract.py |  3 +-
 boa/contracts/vvm/vvm_contract.py  | 78 ++++++++++++++++++++++++------
 boa/interpret.py                   |  4 +-
 4 files changed, 70 insertions(+), 20 deletions(-)

diff --git a/boa/contracts/abi/abi_contract.py b/boa/contracts/abi/abi_contract.py
index 1a357fc4..70104c7b 100644
--- a/boa/contracts/abi/abi_contract.py
+++ b/boa/contracts/abi/abi_contract.py
@@ -1,6 +1,5 @@
 from collections import defaultdict
 from functools import cached_property
-from pathlib import Path
 from typing import Any, Optional, Union
 from warnings import warn
 
@@ -236,7 +235,7 @@ def __init__(
         abi: list[dict],
         functions: list[ABIFunction],
         address: Address,
-        filename: str | Path | None = None,
+        filename: Optional[str] = None,
         env=None,
     ):
         super().__init__(name, env, filename=filename, address=address)
@@ -342,7 +341,7 @@ class ABIContractFactory:
     do any contract deployment.
     """
 
-    def __init__(self, name: str, abi: list[dict], filename: str | Path | None = None):
+    def __init__(self, name: str, abi: list[dict], filename: Optional[str] = None):
         self._name = name
         self._abi = abi
         self.filename = filename
diff --git a/boa/contracts/base_evm_contract.py b/boa/contracts/base_evm_contract.py
index f239a488..5e2a45bf 100644
--- a/boa/contracts/base_evm_contract.py
+++ b/boa/contracts/base_evm_contract.py
@@ -1,5 +1,4 @@
 from dataclasses import dataclass
-from pathlib import Path
 from typing import TYPE_CHECKING, Optional
 
 from eth.abc import ComputationAPI
@@ -26,7 +25,7 @@ def __init__(
         self,
         name: str,
         env: Optional[Env] = None,
-        filename: str | Path | None = None,
+        filename: Optional[str] = None,
         address: Optional[Address] = None,
     ):
         self.env = env or Env.get_singleton()
diff --git a/boa/contracts/vvm/vvm_contract.py b/boa/contracts/vvm/vvm_contract.py
index 37c064df..bc9bc258 100644
--- a/boa/contracts/vvm/vvm_contract.py
+++ b/boa/contracts/vvm/vvm_contract.py
@@ -18,7 +18,7 @@
 VERSION_RE = re.compile(r"\s*#\s*(pragma\s+version|@version)\s+(\d+\.\d+\.\d+)")
 
 
-# TODO: maybe move this up to VVM?
+# TODO: maybe move this up to vvm?
 def _detect_version(source_code: str):
     res = VERSION_RE.findall(source_code)
     if len(res) < 1:
@@ -28,13 +28,17 @@ def _detect_version(source_code: str):
 
 
 class VVMDeployer(ABIContractFactory):
+    """
+    A factory which can be used to create a new contract deployed with vvm.
+    """
+
     def __init__(
         self,
         name: str,
         compiler_output: dict,
         source_code: str,
         vyper_version: str,
-        filename: str | Path | None = None,
+        filename: str | None = None,
     ):
         super().__init__(name, compiler_output["abi"], filename)
         self.compiler_output = compiler_output
@@ -45,17 +49,13 @@ def __init__(
     def bytecode(self):
         return to_bytes(self.compiler_output["bytecode"])
 
-    @property
-    def layout(self):
-        return self.compiler_output["layout"]
-
     @classmethod
     def from_compiler_output(
         cls,
         compiler_output: dict,
         source_code: str,
         vyper_version: str,
-        filename: str | Path | None = None,
+        filename: str | None = None,
         name: str | None = None,
     ):
         if name is None:
@@ -106,6 +106,10 @@ def at(self, address: Address | str) -> "VVMContract":
 
 
 class VVMContract(ABIContract):
+    """
+    A deployed contract compiled with vvm, which is called via ABI.
+    """
+
     def __init__(self, compiler_output, source_code, vyper_version, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.compiler_output = compiler_output
@@ -121,10 +125,25 @@ def bytecode_runtime(self):
         return to_bytes(self.compiler_output["bytecode_runtime"])
 
     def eval(self, code, return_type=None):
+        """
+        Evaluate a vyper statement in the context of this contract.
+        Note that the return_type is necessary to correctly decode the result.
+        WARNING: This is different from the vyper eval() function, which is able
+        to automatically detect the return type.
+        :param code: A vyper statement.
+        :param return_type: The return type of the statement evaluation.
+        :returns: The result of the statement evaluation.
+        """
         return VVMEval(code, self, return_type)()
 
     @cached_property
     def _storage(self):
+        """
+        Allows access to the storage variables of the contract.
+        Note that this is quite slow, as it requires the complete contract to be
+        recompiled.
+        """
+
         def storage():
             return None
 
@@ -134,6 +153,12 @@ def storage():
 
     @cached_property
     def internal(self):
+        """
+        Allows access to internal functions of the contract.
+        Note that this is quite slow, as it requires the complete contract to be
+        recompiled.
+        """
+
         def internal():
             return None
 
@@ -144,16 +169,20 @@ def internal():
                 setattr(internal, function.name, function)
         return internal
 
-    def _compile_metadata_fn_info(self):
-        # todo: move this to vvm
+    def _compile_metadata_fn_info(self) -> dict:
+        """Compiles the metadata for the contract"""
+        # todo: move this to vvm?
         if self.filename is not None:
-            return self._call_vyper(self.filename)
+            return self._get_metadata_from_vyper_executable(self.filename)
         with NamedTemporaryFile(suffix=".vy") as f:
             f.write(self.source_code.encode())
             f.flush()
-            return self._call_vyper(f.name)
+            return self._get_metadata_from_vyper_executable(f.name)
 
-    def _call_vyper(self, filename):
+    def _get_metadata_from_vyper_executable(self, filename: str) -> dict:
+        """
+        Calls the vvm to get the metadata for the contract.
+        """
         stdoutdata, stderrdata, command, proc = vyper_wrapper(
             vyper_binary=get_executable(self.vyper_version),
             f="metadata",
@@ -165,11 +194,14 @@ def _call_vyper(self, filename):
 class _VVMInternal(ABIFunction):
     """
     An ABI function that temporarily changes the bytecode at the contract's address.
+    Subclasses of this class are used to inject code into the contract via the
+    `source_code` property using the vvm, temporarily changing the bytecode
+    at the contract's address.
     """
 
     @cached_property
     def _override_bytecode(self) -> bytes:
-        assert isinstance(self.contract, VVMContract)
+        assert isinstance(self.contract, VVMContract)  # help mypy
         source = "\n".join((self.contract.source_code, self.source_code))
         compiled = vvm.compile_source(source, vyper_version=self.contract.vyper_version)
         return to_bytes(compiled["<stdin>"]["bytecode_runtime"])
@@ -180,7 +212,7 @@ def source_code(self):
 
     def __call__(self, *args, **kwargs):
         env = self.contract.env
-        assert isinstance(self.contract, VVMContract)
+        assert isinstance(self.contract, VVMContract)  # help mypy
         balance_before = env.get_balance(env.eoa)
         env.set_code(self.contract.address, self._override_bytecode)
         env.set_balance(env.eoa, 10**20)
@@ -192,6 +224,11 @@ def __call__(self, *args, **kwargs):
 
 
 class VVMInternalFunction(_VVMInternal):
+    """
+    An internal function that is made available via the `internal` namespace.
+    It will temporarily change the bytecode at the contract's address.
+    """
+
     def __init__(self, meta: dict, contract: VVMContract):
         abi = {
             "anonymous": False,
@@ -238,6 +275,11 @@ def __boa_internal_{self.name}__({fn_sig}){return_sig}:
 
 
 class VVMStorageVariable(_VVMInternal):
+    """
+    A storage variable that is made available via the `storage` namespace.
+    It will temporarily change the bytecode at the contract's address.
+    """
+
     def __init__(self, name, spec, contract):
         abi = {
             "anonymous": False,
@@ -277,6 +319,14 @@ def __boa_private_{self.name}__({args_signature}) -> {self.return_type[0]}:
 
 
 class VVMEval(_VVMInternal):
+    """
+    A Vyper eval statement which can be used to evaluate vyper statements
+    via vvm-compiled contracts. This implementation has some drawbacks:
+    - It is very slow, as it requires the complete contract to be recompiled.
+    - It does not detect the return type, as it is currently not possible.
+    - It will temporarily change the bytecode at the contract's address.
+    """
+
     def __init__(self, code: str, contract: VVMContract, return_type: str = None):
         abi = {
             "anonymous": False,
diff --git a/boa/interpret.py b/boa/interpret.py
index 7db41148..f8b97893 100644
--- a/boa/interpret.py
+++ b/boa/interpret.py
@@ -241,8 +241,10 @@ def _loads_partial_vvm(
     vvm.install_vyper(version=version)
     # TODO: implement caching
     compiled_src = vvm.compile_source(source_code, vyper_version=version)
-
     compiler_output = compiled_src["<stdin>"]
+
+    if filename is not None:
+        filename = str(filename)
     return VVMDeployer.from_compiler_output(
         compiler_output, source_code, version, filename, name
     )

From 1114804dd8fa28db933bad1b1060170c835e9e0b Mon Sep 17 00:00:00 2001
From: Daniel Schiavini <daniel.schiavini@gmail.com>
Date: Tue, 3 Sep 2024 11:03:09 +0200
Subject: [PATCH 05/20] Inline method

---
 boa/contracts/vvm/vvm_contract.py | 30 ++++++++++++++----------------
 1 file changed, 14 insertions(+), 16 deletions(-)

diff --git a/boa/contracts/vvm/vvm_contract.py b/boa/contracts/vvm/vvm_contract.py
index bc9bc258..e50efb21 100644
--- a/boa/contracts/vvm/vvm_contract.py
+++ b/boa/contracts/vvm/vvm_contract.py
@@ -162,33 +162,31 @@ def internal(self):
         def internal():
             return None
 
-        result = self._compile_metadata_fn_info()
+        result = self._compile_function_metadata()
         for fn_name, meta in result.items():
             if meta["visibility"] == "internal":
                 function = VVMInternalFunction(meta, self)
                 setattr(internal, function.name, function)
         return internal
 
-    def _compile_metadata_fn_info(self) -> dict:
-        """Compiles the metadata for the contract"""
+    def _compile_function_metadata(self) -> dict:
+        """Compiles the contract and returns the function metadata"""
+
         # todo: move this to vvm?
+        def run_vyper(filename: str) -> dict:
+            stdoutdata, stderrdata, command, proc = vyper_wrapper(
+                vyper_binary=get_executable(self.vyper_version),
+                f="metadata",
+                source_files=[filename],
+            )
+            return json.loads(stdoutdata)["function_info"]
+
         if self.filename is not None:
-            return self._get_metadata_from_vyper_executable(self.filename)
+            return run_vyper(self.filename)
         with NamedTemporaryFile(suffix=".vy") as f:
             f.write(self.source_code.encode())
             f.flush()
-            return self._get_metadata_from_vyper_executable(f.name)
-
-    def _get_metadata_from_vyper_executable(self, filename: str) -> dict:
-        """
-        Calls the vvm to get the metadata for the contract.
-        """
-        stdoutdata, stderrdata, command, proc = vyper_wrapper(
-            vyper_binary=get_executable(self.vyper_version),
-            f="metadata",
-            source_files=[filename],
-        )
-        return json.loads(stdoutdata)["function_info"]
+            return run_vyper(f.name)
 
 
 class _VVMInternal(ABIFunction):

From 741811678c27eb52998c560c88ad49043440c276 Mon Sep 17 00:00:00 2001
From: Daniel Schiavini <daniel.schiavini@gmail.com>
Date: Fri, 6 Sep 2024 10:18:54 +0200
Subject: [PATCH 06/20] Use the new vvm version

depends on https://github.com/vyperlang/vvm/pull/21
---
 boa/contracts/vvm/vvm_contract.py | 40 +++----------------------------
 boa/interpret.py                  |  4 ++--
 pyproject.toml                    |  2 +-
 3 files changed, 6 insertions(+), 40 deletions(-)

diff --git a/boa/contracts/vvm/vvm_contract.py b/boa/contracts/vvm/vvm_contract.py
index e50efb21..ecf5c02d 100644
--- a/boa/contracts/vvm/vvm_contract.py
+++ b/boa/contracts/vvm/vvm_contract.py
@@ -1,12 +1,7 @@
-import json
-import re
 from functools import cached_property
 from pathlib import Path
-from tempfile import NamedTemporaryFile
 
 import vvm
-from vvm.install import get_executable
-from vvm.wrapper import vyper_wrapper
 from vyper.utils import method_id
 
 from boa.contracts.abi.abi_contract import ABIContract, ABIContractFactory, ABIFunction
@@ -14,18 +9,6 @@
 from boa.rpc import to_bytes
 from boa.util.abi import Address
 
-# TODO: maybe this doesn't detect release candidates
-VERSION_RE = re.compile(r"\s*#\s*(pragma\s+version|@version)\s+(\d+\.\d+\.\d+)")
-
-
-# TODO: maybe move this up to vvm?
-def _detect_version(source_code: str):
-    res = VERSION_RE.findall(source_code)
-    if len(res) < 1:
-        return None
-    # TODO: handle len(res) > 1
-    return res[0][1]
-
 
 class VVMDeployer(ABIContractFactory):
     """
@@ -162,32 +145,15 @@ def internal(self):
         def internal():
             return None
 
-        result = self._compile_function_metadata()
+        result = vvm.compile_source(
+            self.source_code, vyper_version=self.vyper_version, output_format="metadata"
+        )["function_info"]
         for fn_name, meta in result.items():
             if meta["visibility"] == "internal":
                 function = VVMInternalFunction(meta, self)
                 setattr(internal, function.name, function)
         return internal
 
-    def _compile_function_metadata(self) -> dict:
-        """Compiles the contract and returns the function metadata"""
-
-        # todo: move this to vvm?
-        def run_vyper(filename: str) -> dict:
-            stdoutdata, stderrdata, command, proc = vyper_wrapper(
-                vyper_binary=get_executable(self.vyper_version),
-                f="metadata",
-                source_files=[filename],
-            )
-            return json.loads(stdoutdata)["function_info"]
-
-        if self.filename is not None:
-            return run_vyper(self.filename)
-        with NamedTemporaryFile(suffix=".vy") as f:
-            f.write(self.source_code.encode())
-            f.flush()
-            return run_vyper(f.name)
-
 
 class _VVMInternal(ABIFunction):
     """
diff --git a/boa/interpret.py b/boa/interpret.py
index f8b97893..1cea7259 100644
--- a/boa/interpret.py
+++ b/boa/interpret.py
@@ -21,7 +21,7 @@
 from vyper.utils import sha256sum
 
 from boa.contracts.abi.abi_contract import ABIContractFactory
-from boa.contracts.vvm.vvm_contract import VVMDeployer, _detect_version
+from boa.contracts.vvm.vvm_contract import VVMDeployer
 from boa.contracts.vyper.vyper_contract import (
     VyperBlueprint,
     VyperContract,
@@ -212,7 +212,7 @@ def loads_partial(
     if dedent:
         source_code = textwrap.dedent(source_code)
 
-    version = _detect_version(source_code)
+    version = vvm.detect_vyper_version_from_source(source_code)
     if version is not None and version != vyper.__version__:
         return _loads_partial_vvm(source_code, version, filename, name)
 
diff --git a/pyproject.toml b/pyproject.toml
index 12ab92db..a7bf1cb9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,7 +25,7 @@ dependencies = [
     "pytest-cov",
 
     # required to compile older versions of vyper
-    "vvm",
+    "vvm>=0.2.2",
 
     # eth-rlp requirement, not installed by default with 3.12
     "typing-extensions",

From 9a0cd05fe4b2e948a319ffeb81afab5d2000ed0f Mon Sep 17 00:00:00 2001
From: Daniel Schiavini <daniel.schiavini@gmail.com>
Date: Wed, 11 Sep 2024 11:41:34 +0200
Subject: [PATCH 07/20] Recursion, review comments

---
 CONTRIBUTING.md                          |  2 +-
 boa/contracts/vvm/vvm_contract.py        | 30 ++++++++++++++----------
 tests/unitary/contracts/vvm/mock_3_10.vy |  4 ++--
 tests/unitary/contracts/vvm/test_vvm.py  |  6 ++---
 4 files changed, 23 insertions(+), 19 deletions(-)

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index ab445856..b2dba295 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -53,7 +53,7 @@ source venv/bin/activate
 # Install dev requirements
 pip install -r dev-requirements.txt
 # Install prod requirements (in the pyproject.tom)
-pip install . 
+pip install .
 ```
 
 *Note: When you delete your terminal/shell, you will need to reactivate this virtual environment again each time. To exit this python virtual environment, type `deactivate`*
diff --git a/boa/contracts/vvm/vvm_contract.py b/boa/contracts/vvm/vvm_contract.py
index ecf5c02d..5275a80d 100644
--- a/boa/contracts/vvm/vvm_contract.py
+++ b/boa/contracts/vvm/vvm_contract.py
@@ -1,3 +1,4 @@
+import re
 from functools import cached_property
 from pathlib import Path
 
@@ -171,19 +172,20 @@ def _override_bytecode(self) -> bytes:
         return to_bytes(compiled["<stdin>"]["bytecode_runtime"])
 
     @property
-    def source_code(self):
-        raise NotImplementedError  # to be implemented in subclasses
+    def source_code(self) -> str:
+        """
+        Returns the source code an internal function.
+        Must be implemented in subclasses.
+        """
+        raise NotImplementedError
 
     def __call__(self, *args, **kwargs):
         env = self.contract.env
         assert isinstance(self.contract, VVMContract)  # help mypy
-        balance_before = env.get_balance(env.eoa)
         env.set_code(self.contract.address, self._override_bytecode)
-        env.set_balance(env.eoa, 10**20)
         try:
             return super().__call__(*args, **kwargs)
         finally:
-            env.set_balance(env.eoa, balance_before)
             env.set_code(self.contract.address, self.contract.bytecode_runtime)
 
 
@@ -245,19 +247,21 @@ class VVMStorageVariable(_VVMInternal):
     """
 
     def __init__(self, name, spec, contract):
+        value_type = spec["type"]
+        inputs = []
+        regex = re.compile(r"^HashMap\[([^[]+), (.+)]$")
+
+        while value_type.startswith("HashMap"):
+            key_type, value_type = regex.match(value_type).groups()
+            inputs.append({"name": f"key{len(inputs)}", "type": key_type})
+
         abi = {
             "anonymous": False,
-            "inputs": [],
-            "outputs": [{"name": name, "type": spec["type"]}],
+            "inputs": inputs,
+            "outputs": [{"name": name, "type": value_type}],
             "name": name,
             "type": "function",
         }
-
-        if spec["type"].startswith("HashMap"):
-            key_type, value_type = spec["type"][8:-1].split(",")
-            abi["inputs"] = [{"name": "key", "type": key_type}]
-            abi["outputs"] = [{"name": "value", "type": value_type.strip()}]
-
         super().__init__(abi, contract.contract_name)
         self.contract = contract
 
diff --git a/tests/unitary/contracts/vvm/mock_3_10.vy b/tests/unitary/contracts/vvm/mock_3_10.vy
index 339f1ce9..7d5451bc 100644
--- a/tests/unitary/contracts/vvm/mock_3_10.vy
+++ b/tests/unitary/contracts/vvm/mock_3_10.vy
@@ -3,7 +3,7 @@
 foo: public(uint256)
 bar: public(uint256)
 
-map: HashMap[address, uint256]
+map: HashMap[address, HashMap[uint8, uint256]]
 is_empty: bool
 
 @external
@@ -18,5 +18,5 @@ def set_map(x: uint256):
 
 @internal
 def _set_map(addr: address, x: uint256):
-    self.map[addr] = x
+    self.map[addr][0] = x
     self.is_empty = False
diff --git a/tests/unitary/contracts/vvm/test_vvm.py b/tests/unitary/contracts/vvm/test_vvm.py
index 9e29b349..999a000b 100644
--- a/tests/unitary/contracts/vvm/test_vvm.py
+++ b/tests/unitary/contracts/vvm/test_vvm.py
@@ -38,10 +38,10 @@ def test_loads_vvm():
 def test_vvm_storage():
     contract = boa.loads(mock_3_10_code, 43)
     assert contract._storage.is_empty.get()
-    assert contract._storage.map.get(boa.env.eoa) == 0
+    assert contract._storage.map.get(boa.env.eoa, 0) == 0
     contract.set_map(69)
     assert not contract._storage.is_empty.get()
-    assert contract._storage.map.get(boa.env.eoa) == 69
+    assert contract._storage.map.get(boa.env.eoa, 0) == 69
 
 
 def test_vvm_internal():
@@ -49,7 +49,7 @@ def test_vvm_internal():
     assert not hasattr(contract.internal, "set_map")
     address = boa.env.generate_address()
     contract.internal._set_map(address, 69)
-    assert contract._storage.map.get(address) == 69
+    assert contract._storage.map.get(address, 0) == 69
 
 
 def test_vvm_eval():

From b375574173ab49913b9543fefa147c74a33dcac0 Mon Sep 17 00:00:00 2001
From: Daniel Schiavini <daniel.schiavini@gmail.com>
Date: Wed, 18 Sep 2024 20:32:04 +0200
Subject: [PATCH 08/20] Use vvm from https://github.com/vyperlang/vvm/pull/26

---
 boa/interpret.py                          | 11 ++++++++---
 pyproject.toml                            |  5 +++--
 tests/unitary/fixtures/module_contract.vy |  2 +-
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/boa/interpret.py b/boa/interpret.py
index 9675bb2f..9d1f9a0d 100644
--- a/boa/interpret.py
+++ b/boa/interpret.py
@@ -8,6 +8,7 @@
 
 import vvm
 import vyper
+from vvm.exceptions import UnexpectedVersionError
 from vyper.cli.vyper_compile import get_search_paths
 from vyper.compiler.input_bundle import (
     ABIInput,
@@ -212,9 +213,13 @@ def loads_partial(
     if dedent:
         source_code = textwrap.dedent(source_code)
 
-    version = vvm.detect_vyper_version_from_source(source_code)
-    if version is not None and version != vyper.__version__:
-        return _loads_partial_vvm(source_code, version, filename)
+    try:
+        version = vvm.detect_vyper_version_from_source(source_code)
+        if str(version) != vyper.__version__:
+            return _loads_partial_vvm(source_code, version, filename)
+    except UnexpectedVersionError as e:
+        if e.args[0] != "No version detected in source code":
+            raise
 
     compiler_args = compiler_args or {}
 
diff --git a/pyproject.toml b/pyproject.toml
index fe86c392..3c79b22e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -15,7 +15,7 @@ classifiers = ["Topic :: Software Development"]
 
 # Requirements
 dependencies = [
-    "vyper>=0.4.0",
+    "vyper>=0.4.0,<0.5.0",
     "eth-stdlib>=0.2.7,<0.3.0",
     "eth-abi",
     "py-evm>=0.10.0b4",
@@ -25,7 +25,8 @@ dependencies = [
     "pytest-cov",
 
     # required to compile older versions of vyper
-    "vvm>=0.2.2",
+#    "vvm>=0.3.1,<0.4.0",
+    "vvm @ git+https://github.com/DanielSchiavini/vvm.git@dependencies",
 
     # eth-rlp requirement, not installed by default with 3.12
     "typing-extensions",
diff --git a/tests/unitary/fixtures/module_contract.vy b/tests/unitary/fixtures/module_contract.vy
index 886424db..45422642 100644
--- a/tests/unitary/fixtures/module_contract.vy
+++ b/tests/unitary/fixtures/module_contract.vy
@@ -1,4 +1,4 @@
-# pragma version ^0.4.0
+# pragma version >=0.4.0
 
 import module_lib
 

From 8b778dd75ce963c84936911b06b0ef71ec2dd457 Mon Sep 17 00:00:00 2001
From: Daniel Schiavini <daniel.schiavini@gmail.com>
Date: Mon, 23 Sep 2024 09:51:40 +0200
Subject: [PATCH 09/20] Update vvm

---
 boa/interpret.py | 14 +++++---------
 pyproject.toml   |  3 +--
 2 files changed, 6 insertions(+), 11 deletions(-)

diff --git a/boa/interpret.py b/boa/interpret.py
index f35c74cb..cf0fc71b 100644
--- a/boa/interpret.py
+++ b/boa/interpret.py
@@ -8,7 +8,7 @@
 
 import vvm
 import vyper
-from vvm.exceptions import UnexpectedVersionError
+from packaging.version import Version
 from vyper.cli.vyper_compile import get_search_paths
 from vyper.compiler.input_bundle import (
     ABIInput,
@@ -214,14 +214,10 @@ def loads_partial(
     if dedent:
         source_code = textwrap.dedent(source_code)
 
-    try:
-        version = vvm.detect_vyper_version_from_source(source_code)
-        if str(version) != vyper.__version__:
-            # TODO: pass name to loads_partial_vvm, not filename
-            return _loads_partial_vvm(source_code, version, filename)
-    except UnexpectedVersionError as e:
-        if e.args[0] != "No version detected in source code":
-            raise
+    version = vvm.detect_vyper_version_from_source(source_code)
+    if version is not None and version != Version(vyper.__version__):
+        # TODO: pass name to loads_partial_vvm, not filename
+        return _loads_partial_vvm(source_code, str(version), filename)
 
     compiler_args = compiler_args or {}
 
diff --git a/pyproject.toml b/pyproject.toml
index 3c79b22e..b7a93bf9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,8 +25,7 @@ dependencies = [
     "pytest-cov",
 
     # required to compile older versions of vyper
-#    "vvm>=0.3.1,<0.4.0",
-    "vvm @ git+https://github.com/DanielSchiavini/vvm.git@dependencies",
+    "vvm>=0.3.1,<0.4.0",
 
     # eth-rlp requirement, not installed by default with 3.12
     "typing-extensions",

From eb4518f45dfc71d65be150e185a1a7aa0d7376bc Mon Sep 17 00:00:00 2001
From: Daniel Schiavini <daniel.schiavini@gmail.com>
Date: Mon, 23 Sep 2024 12:46:23 +0200
Subject: [PATCH 10/20] Extract regex

---
 boa/contracts/vvm/vvm_contract.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/boa/contracts/vvm/vvm_contract.py b/boa/contracts/vvm/vvm_contract.py
index 5275a80d..3df6c8d1 100644
--- a/boa/contracts/vvm/vvm_contract.py
+++ b/boa/contracts/vvm/vvm_contract.py
@@ -246,13 +246,14 @@ class VVMStorageVariable(_VVMInternal):
     It will temporarily change the bytecode at the contract's address.
     """
 
+    _hashmap_regex = re.compile(r"^HashMap\[([^[]+), (.+)]$")
+
     def __init__(self, name, spec, contract):
         value_type = spec["type"]
         inputs = []
-        regex = re.compile(r"^HashMap\[([^[]+), (.+)]$")
 
         while value_type.startswith("HashMap"):
-            key_type, value_type = regex.match(value_type).groups()
+            key_type, value_type = self._hashmap_regex.match(value_type).groups()
             inputs.append({"name": f"key{len(inputs)}", "type": key_type})
 
         abi = {

From 8de1b7f678d99077c9b62b216ca3b70012de505e Mon Sep 17 00:00:00 2001
From: Daniel Schiavini <daniel.schiavini@gmail.com>
Date: Thu, 26 Sep 2024 20:33:50 +0200
Subject: [PATCH 11/20] Review comments

---
 boa/contracts/vvm/vvm_contract.py         | 7 ++++---
 boa/interpret.py                          | 6 +++---
 pyproject.toml                            | 4 ++--
 tests/unitary/contracts/vvm/mock_3_10.vy  | 4 ++--
 tests/unitary/contracts/vvm/test_vvm.py   | 6 +++---
 tests/unitary/fixtures/module_contract.vy | 2 +-
 6 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/boa/contracts/vvm/vvm_contract.py b/boa/contracts/vvm/vvm_contract.py
index 3df6c8d1..f1a97963 100644
--- a/boa/contracts/vvm/vvm_contract.py
+++ b/boa/contracts/vvm/vvm_contract.py
@@ -1,6 +1,7 @@
 import re
 from functools import cached_property
 from pathlib import Path
+from typing import Optional
 
 import vvm
 from vyper.utils import method_id
@@ -22,7 +23,7 @@ def __init__(
         compiler_output: dict,
         source_code: str,
         vyper_version: str,
-        filename: str | None = None,
+        filename: Optional[str] = None,
     ):
         super().__init__(name, compiler_output["abi"], filename)
         self.compiler_output = compiler_output
@@ -39,8 +40,8 @@ def from_compiler_output(
         compiler_output: dict,
         source_code: str,
         vyper_version: str,
-        filename: str | None = None,
-        name: str | None = None,
+        filename: Optional[str] = None,
+        name: Optional[str] = None,
     ):
         if name is None:
             name = Path(filename).stem if filename is not None else "<VVMContract>"
diff --git a/boa/interpret.py b/boa/interpret.py
index cf0fc71b..1464e859 100644
--- a/boa/interpret.py
+++ b/boa/interpret.py
@@ -4,7 +4,7 @@
 from importlib.machinery import SourceFileLoader
 from importlib.util import spec_from_loader
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 import vvm
 import vyper
@@ -237,8 +237,8 @@ def load_partial(filename: str, compiler_args=None):
 def _loads_partial_vvm(
     source_code: str,
     version: str,
-    filename: str | Path | None = None,
-    name: str | None = None,
+    filename: Optional[str | Path] = None,
+    name: Optional[str] = None,
 ):
     global _disk_cache
     if filename is not None:
diff --git a/pyproject.toml b/pyproject.toml
index b7a93bf9..17a0b95a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -15,7 +15,7 @@ classifiers = ["Topic :: Software Development"]
 
 # Requirements
 dependencies = [
-    "vyper>=0.4.0,<0.5.0",
+    "vyper>=0.4.0",
     "eth-stdlib>=0.2.7,<0.3.0",
     "eth-abi",
     "py-evm>=0.10.0b4",
@@ -25,7 +25,7 @@ dependencies = [
     "pytest-cov",
 
     # required to compile older versions of vyper
-    "vvm>=0.3.1,<0.4.0",
+    "vvm>=0.3.1",
 
     # eth-rlp requirement, not installed by default with 3.12
     "typing-extensions",
diff --git a/tests/unitary/contracts/vvm/mock_3_10.vy b/tests/unitary/contracts/vvm/mock_3_10.vy
index 7d5451bc..28e5e909 100644
--- a/tests/unitary/contracts/vvm/mock_3_10.vy
+++ b/tests/unitary/contracts/vvm/mock_3_10.vy
@@ -3,7 +3,7 @@
 foo: public(uint256)
 bar: public(uint256)
 
-map: HashMap[address, HashMap[uint8, uint256]]
+hash_map: HashMap[address, HashMap[uint8, uint256]]
 is_empty: bool
 
 @external
@@ -18,5 +18,5 @@ def set_map(x: uint256):
 
 @internal
 def _set_map(addr: address, x: uint256):
-    self.map[addr][0] = x
+    self.hash_map[addr][0] = x
     self.is_empty = False
diff --git a/tests/unitary/contracts/vvm/test_vvm.py b/tests/unitary/contracts/vvm/test_vvm.py
index 999a000b..19a2a56a 100644
--- a/tests/unitary/contracts/vvm/test_vvm.py
+++ b/tests/unitary/contracts/vvm/test_vvm.py
@@ -38,10 +38,10 @@ def test_loads_vvm():
 def test_vvm_storage():
     contract = boa.loads(mock_3_10_code, 43)
     assert contract._storage.is_empty.get()
-    assert contract._storage.map.get(boa.env.eoa, 0) == 0
+    assert contract._storage.hash_map.get(boa.env.eoa, 0) == 0
     contract.set_map(69)
     assert not contract._storage.is_empty.get()
-    assert contract._storage.map.get(boa.env.eoa, 0) == 69
+    assert contract._storage.hash_map.get(boa.env.eoa, 0) == 69
 
 
 def test_vvm_internal():
@@ -49,7 +49,7 @@ def test_vvm_internal():
     assert not hasattr(contract.internal, "set_map")
     address = boa.env.generate_address()
     contract.internal._set_map(address, 69)
-    assert contract._storage.map.get(address, 0) == 69
+    assert contract._storage.hash_map.get(address, 0) == 69
 
 
 def test_vvm_eval():
diff --git a/tests/unitary/fixtures/module_contract.vy b/tests/unitary/fixtures/module_contract.vy
index 45422642..9a243065 100644
--- a/tests/unitary/fixtures/module_contract.vy
+++ b/tests/unitary/fixtures/module_contract.vy
@@ -1,4 +1,4 @@
-# pragma version >=0.4.0
+# pragma version ~=0.4.0
 
 import module_lib
 

From e101d0b0ad5033ee8e4e3f82f5bf00ec2c0b066d Mon Sep 17 00:00:00 2001
From: Daniel Schiavini <daniel.schiavini@gmail.com>
Date: Fri, 4 Oct 2024 11:04:28 +0200
Subject: [PATCH 12/20] refactor: extract function

---
 boa/contracts/vvm/vvm_contract.py | 29 +++++++++++++++++++----------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/boa/contracts/vvm/vvm_contract.py b/boa/contracts/vvm/vvm_contract.py
index c1605402..e440890f 100644
--- a/boa/contracts/vvm/vvm_contract.py
+++ b/boa/contracts/vvm/vvm_contract.py
@@ -284,20 +284,12 @@ class VVMStorageVariable(_VVMInternal):
     It will temporarily change the bytecode at the contract's address.
     """
 
-    _hashmap_regex = re.compile(r"^HashMap\[([^[]+), (.+)]$")
-
     def __init__(self, name, spec, contract):
-        value_type = spec["type"]
-        inputs = []
-
-        while value_type.startswith("HashMap"):
-            key_type, value_type = self._hashmap_regex.match(value_type).groups()
-            inputs.append({"name": f"key{len(inputs)}", "type": key_type})
-
+        inputs, output_type = _get_storage_variable_types(spec)
         abi = {
             "anonymous": False,
             "inputs": inputs,
-            "outputs": [{"name": name, "type": value_type}],
+            "outputs": [{"name": name, "type": output_type}],
             "name": name,
             "type": "function",
         }
@@ -359,3 +351,20 @@ def source_code(self):
 def __boa_debug__() {return_sig}:
     {debug_body}
 """
+
+
+def _get_storage_variable_types(spec: dict) -> tuple[list[dict], str]:
+    """
+    Get the types of a storage variable
+    :param spec: The storage variable specification.
+    :return: The types of the storage variable:
+    1. A list of dictionaries containing the input types.
+    2. The output type name.
+    """
+    hashmap_regex = re.compile(r"^HashMap\[([^[]+), (.+)]$")
+    output_type = spec["type"]
+    inputs: list[dict] = []
+    while output_type.startswith("HashMap"):
+        key_type, output_type = hashmap_regex.match(output_type).groups()  # type: ignore
+        inputs.append({"name": f"key{len(inputs)}", "type": key_type})
+    return inputs, output_type

From dcfe94156dd4f3ff3c9c4056b637d05572b325cf Mon Sep 17 00:00:00 2001
From: Daniel Schiavini <daniel.schiavini@gmail.com>
Date: Fri, 4 Oct 2024 11:20:26 +0200
Subject: [PATCH 13/20] feat: cache all vvm compile calls

---
 boa/contracts/vvm/vvm_contract.py |  8 +++--
 boa/interpret.py                  | 59 +++++++------------------------
 boa/util/cached_vvm.py            | 26 ++++++++++++++
 boa/util/disk_cache.py            | 36 +++++++++++++++++++
 tests/unitary/utils/test_cache.py |  9 ++---
 5 files changed, 84 insertions(+), 54 deletions(-)
 create mode 100644 boa/util/cached_vvm.py

diff --git a/boa/contracts/vvm/vvm_contract.py b/boa/contracts/vvm/vvm_contract.py
index e440890f..45b94c36 100644
--- a/boa/contracts/vvm/vvm_contract.py
+++ b/boa/contracts/vvm/vvm_contract.py
@@ -3,12 +3,12 @@
 from pathlib import Path
 from typing import Optional
 
-import vvm
 from vyper.utils import method_id
 
 from boa.contracts.abi.abi_contract import ABIContract, ABIContractFactory, ABIFunction
 from boa.environment import Env
 from boa.rpc import to_bytes
+from boa.util import cached_vvm
 from boa.util.abi import Address
 from boa.util.eip5202 import generate_blueprint_bytecode
 
@@ -184,7 +184,7 @@ def internal(self):
         def internal():
             return None
 
-        result = vvm.compile_source(
+        result = cached_vvm.compile_source(
             self.source_code, vyper_version=self.vyper_version, output_format="metadata"
         )["function_info"]
         for fn_name, meta in result.items():
@@ -206,7 +206,9 @@ class _VVMInternal(ABIFunction):
     def _override_bytecode(self) -> bytes:
         assert isinstance(self.contract, VVMContract)  # help mypy
         source = "\n".join((self.contract.source_code, self.source_code))
-        compiled = vvm.compile_source(source, vyper_version=self.contract.vyper_version)
+        compiled = cached_vvm.compile_source(
+            source, vyper_version=self.contract.vyper_version
+        )
         return to_bytes(compiled["<stdin>"]["bytecode_runtime"])
 
     @property
diff --git a/boa/interpret.py b/boa/interpret.py
index 1464e859..4374282b 100644
--- a/boa/interpret.py
+++ b/boa/interpret.py
@@ -31,8 +31,9 @@
 from boa.environment import Env
 from boa.explorer import Etherscan, get_etherscan
 from boa.rpc import json
+from boa.util import cached_vvm
 from boa.util.abi import Address
-from boa.util.disk_cache import DiskCache
+from boa.util.disk_cache import get_disk_cache, get_search_path
 
 if TYPE_CHECKING:
     from vyper.semantics.analysis.base import ImportInfo
@@ -40,31 +41,6 @@
 _Contract = Union[VyperContract, VyperBlueprint]
 
 
-_disk_cache = None
-_search_path = None
-
-
-def set_search_path(path: list[str]):
-    global _search_path
-    _search_path = path
-
-
-def set_cache_dir(cache_dir="~/.cache/titanoboa"):
-    global _disk_cache
-    if cache_dir is None:
-        _disk_cache = None
-        return
-    compiler_version = f"{vyper.__version__}.{vyper.__commit__}"
-    _disk_cache = DiskCache(cache_dir, compiler_version)
-
-
-def disable_cache():
-    set_cache_dir(None)
-
-
-set_cache_dir()  # enable caching, by default!
-
-
 class BoaImporter(MetaPathFinder):
     def find_spec(self, fullname, path, target=None):
         path = Path(fullname.replace(".", "/")).with_suffix(".vy")
@@ -130,8 +106,6 @@ def get_module_fingerprint(
 def compiler_data(
     source_code: str, contract_name: str, filename: str | Path, deployer=None, **kwargs
 ) -> CompilerData:
-    global _disk_cache, _search_path
-
     path = Path(contract_name)
     resolved_path = Path(filename).resolve(strict=False)
 
@@ -139,12 +113,14 @@ def compiler_data(
         contents=source_code, source_id=-1, path=path, resolved_path=resolved_path
     )
 
-    search_paths = get_search_paths(_search_path)
+    search_paths = get_search_paths(get_search_path())
     input_bundle = FilesystemInputBundle(search_paths)
 
     settings = Settings(**kwargs)
     ret = CompilerData(file_input, input_bundle, settings)
-    if _disk_cache is None:
+
+    disk_cache = get_disk_cache()
+    if disk_cache is None:
         return ret
 
     with anchor_settings(ret.settings):
@@ -164,7 +140,7 @@ def get_compiler_data():
     assert isinstance(deployer, type) or deployer is None
     deployer_id = repr(deployer)  # a unique str identifying the deployer class
     cache_key = str((contract_name, fingerprint, kwargs, deployer_id))
-    return _disk_cache.caching_lookup(cache_key, get_compiler_data)
+    return disk_cache.caching_lookup(cache_key, get_compiler_data)
 
 
 def load(filename: str | Path, *args, **kwargs) -> _Contract:  # type: ignore
@@ -240,28 +216,17 @@ def _loads_partial_vvm(
     filename: Optional[str | Path] = None,
     name: Optional[str] = None,
 ):
-    global _disk_cache
     if filename is not None:
         filename = str(filename)
 
     # install the requested version if not already installed
     vvm.install_vyper(version=version)
 
-    def _compile():
-        compiled_src = vvm.compile_source(source_code, vyper_version=version)
-        compiler_output = compiled_src["<stdin>"]
-        return VVMDeployer.from_compiler_output(
-            compiler_output, source_code, version, filename, name
-        )
-
-    # Ensure the cache is initialized
-    if _disk_cache is None:
-        return _compile()
-
-    # Generate a unique cache key
-    cache_key = f"{source_code}:{version}"
-    # Check the cache and return the result if available
-    return _disk_cache.caching_lookup(cache_key, _compile)
+    compiled_src = cached_vvm.compile_source(source_code, vyper_version=version)
+    compiler_output = compiled_src["<stdin>"]
+    return VVMDeployer.from_compiler_output(
+        compiler_output, source_code, version, filename, name
+    )
 
 
 def from_etherscan(
diff --git a/boa/util/cached_vvm.py b/boa/util/cached_vvm.py
new file mode 100644
index 00000000..653543b6
--- /dev/null
+++ b/boa/util/cached_vvm.py
@@ -0,0 +1,26 @@
+from typing import Any
+
+import vvm
+
+from boa.util.disk_cache import get_disk_cache
+
+
+def compile_source(*args, **kwargs) -> Any:
+    """
+    Compile Vyper source code via the VVM.
+    When a disk cache is available, the result of the compilation is cached.
+    Note the cache only works if the function is called the same way (args/kwargs).
+    :param args: Arguments to pass to vvm.compile_source
+    :param kwargs: Keyword arguments to pass to vvm.compile_source
+    :return: Compilation output
+    """
+    disk_cache = get_disk_cache()
+
+    def _compile():
+        return vvm.compile_source(*args, **kwargs)
+
+    if disk_cache is None:
+        return _compile()
+
+    cache_key = f"{args}{kwargs}"
+    return disk_cache.caching_lookup(cache_key, _compile)
diff --git a/boa/util/disk_cache.py b/boa/util/disk_cache.py
index 6912ec51..9598ce5a 100644
--- a/boa/util/disk_cache.py
+++ b/boa/util/disk_cache.py
@@ -5,6 +5,9 @@
 import threading
 import time
 from pathlib import Path
+from typing import Optional
+
+import vyper
 
 _ONE_WEEK = 7 * 24 * 3600
 
@@ -77,3 +80,36 @@ def caching_lookup(self, string, func):
         # because worst case we will just rebuild the item
         tmp_p.rename(p)
         return res
+
+
+_disk_cache = None
+_search_path = None
+
+
+def get_search_path() -> Optional[list[str]]:
+    return _search_path
+
+
+def set_search_path(path: list[str]):
+    global _search_path
+    _search_path = path
+
+
+def get_disk_cache() -> Optional[DiskCache]:
+    return _disk_cache
+
+
+def set_cache_dir(cache_dir: Optional[str] = "~/.cache/titanoboa"):
+    global _disk_cache
+    if cache_dir is None:
+        _disk_cache = None
+        return
+    compiler_version = f"{vyper.__version__}.{vyper.__commit__}"
+    _disk_cache = DiskCache(cache_dir, compiler_version)
+
+
+def disable_cache():
+    set_cache_dir(None)
+
+
+set_cache_dir()  # enable caching, by default!
diff --git a/tests/unitary/utils/test_cache.py b/tests/unitary/utils/test_cache.py
index 64659821..c21cfe84 100644
--- a/tests/unitary/utils/test_cache.py
+++ b/tests/unitary/utils/test_cache.py
@@ -4,12 +4,13 @@
 from vyper.compiler import CompilerData
 
 from boa.contracts.vyper.vyper_contract import VyperDeployer
-from boa.interpret import _disk_cache, _loads_partial_vvm, compiler_data, set_cache_dir
+from boa.interpret import _loads_partial_vvm, compiler_data, get_disk_cache
+from boa.util.disk_cache import set_cache_dir
 
 
 @pytest.fixture(autouse=True)
 def cache_dir(tmp_path):
-    tmp = _disk_cache.cache_dir
+    tmp = get_disk_cache().cache_dir
     try:
         set_cache_dir(tmp_path)
         yield
@@ -21,7 +22,7 @@ def test_cache_contract_name():
     code = """
 x: constant(int128) = 1000
 """
-    assert _disk_cache is not None
+    assert get_disk_cache() is not None
     test1 = compiler_data(code, "test1", __file__, VyperDeployer)
     test2 = compiler_data(code, "test2", __file__, VyperDeployer)
     test3 = compiler_data(code, "test1", __file__, VyperDeployer)
@@ -36,7 +37,7 @@ def test_cache_vvm():
 """
     version = "0.2.8"
     version2 = "0.3.1"
-    assert _disk_cache is not None
+    assert get_disk_cache() is not None
 
     # Mock vvm.compile_source
     with patch("vvm.compile_source") as mock_compile:

From 23ed6ecf4efc68f001f664d2231339074e08e84a Mon Sep 17 00:00:00 2001
From: Daniel Schiavini <daniel.schiavini@gmail.com>
Date: Mon, 7 Oct 2024 16:01:39 +0200
Subject: [PATCH 14/20] fix: revert search path changes

---
 boa/interpret.py       | 13 +++++++++++--
 boa/util/disk_cache.py | 10 ----------
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/boa/interpret.py b/boa/interpret.py
index 4374282b..0ff77392 100644
--- a/boa/interpret.py
+++ b/boa/interpret.py
@@ -33,13 +33,20 @@
 from boa.rpc import json
 from boa.util import cached_vvm
 from boa.util.abi import Address
-from boa.util.disk_cache import get_disk_cache, get_search_path
+from boa.util.disk_cache import get_disk_cache
 
 if TYPE_CHECKING:
     from vyper.semantics.analysis.base import ImportInfo
 
 _Contract = Union[VyperContract, VyperBlueprint]
 
+_search_path = None
+
+
+def set_search_path(path: list[str]):
+    global _search_path
+    _search_path = path
+
 
 class BoaImporter(MetaPathFinder):
     def find_spec(self, fullname, path, target=None):
@@ -106,6 +113,8 @@ def get_module_fingerprint(
 def compiler_data(
     source_code: str, contract_name: str, filename: str | Path, deployer=None, **kwargs
 ) -> CompilerData:
+    global _search_path
+
     path = Path(contract_name)
     resolved_path = Path(filename).resolve(strict=False)
 
@@ -113,7 +122,7 @@ def compiler_data(
         contents=source_code, source_id=-1, path=path, resolved_path=resolved_path
     )
 
-    search_paths = get_search_paths(get_search_path())
+    search_paths = get_search_paths(_search_path)
     input_bundle = FilesystemInputBundle(search_paths)
 
     settings = Settings(**kwargs)
diff --git a/boa/util/disk_cache.py b/boa/util/disk_cache.py
index 9598ce5a..19160f00 100644
--- a/boa/util/disk_cache.py
+++ b/boa/util/disk_cache.py
@@ -83,16 +83,6 @@ def caching_lookup(self, string, func):
 
 
 _disk_cache = None
-_search_path = None
-
-
-def get_search_path() -> Optional[list[str]]:
-    return _search_path
-
-
-def set_search_path(path: list[str]):
-    global _search_path
-    _search_path = path
 
 
 def get_disk_cache() -> Optional[DiskCache]:

From 42eca98cf7f8a7c7e30812e1116977d09d8e0857 Mon Sep 17 00:00:00 2001
From: Daniel Schiavini <daniel.schiavini@gmail.com>
Date: Fri, 18 Oct 2024 10:45:31 +0200
Subject: [PATCH 15/20] feat: implement function injection instead of eval

---
 boa/contracts/vvm/vvm_contract.py       | 62 +++++++++++--------------
 tests/unitary/contracts/vvm/test_vvm.py | 30 ++++++++++--
 2 files changed, 54 insertions(+), 38 deletions(-)

diff --git a/boa/contracts/vvm/vvm_contract.py b/boa/contracts/vvm/vvm_contract.py
index 738805b0..6a136dea 100644
--- a/boa/contracts/vvm/vvm_contract.py
+++ b/boa/contracts/vvm/vvm_contract.py
@@ -146,17 +146,19 @@ def bytecode(self):
     def bytecode_runtime(self):
         return to_bytes(self.compiler_output["bytecode_runtime"])
 
-    def eval(self, code, return_type=None):
+    def inject_function(self, fn_source_code, force=False):
         """
-        Evaluate a vyper statement in the context of this contract.
-        Note that the return_type is necessary to correctly decode the result.
-        WARNING: This is different from the vyper eval() function, which is able
-        to automatically detect the return type.
-        :param code: A vyper statement.
-        :param return_type: The return type of the statement evaluation.
+        Inject a function into this VVM Contract without affecting the
+        contract's source code. useful for testing private functionality.
+        :param fn_source_code: The source code of the function to inject.
+        :param force: If True, the function will be injected even if it already exists.
         :returns: The result of the statement evaluation.
         """
-        return VVMEval(code, self, return_type)()
+        fn = VVMInjectedFunction(fn_source_code, self)
+        if hasattr(self, fn.name) and not force:
+            raise ValueError(f"Function {fn.name} already exists on contract.")
+        setattr(self, fn.name, fn)
+        fn.contract = self
 
     @cached_property
     def _storage(self):
@@ -204,12 +206,16 @@ class _VVMInternal(ABIFunction):
 
     @cached_property
     def _override_bytecode(self) -> bytes:
+        return to_bytes(self._compiler_output["bytecode_runtime"])
+
+    @cached_property
+    def _compiler_output(self):
         assert isinstance(self.contract, VVMContract)  # help mypy
         source = "\n".join((self.contract.source_code, self.source_code))
         compiled = cached_vvm.compile_source(
             source, vyper_version=self.contract.vyper_version
         )
-        return to_bytes(compiled["<stdin>"]["bytecode_runtime"])
+        return compiled["<stdin>"]
 
     @property
     def source_code(self) -> str:
@@ -319,40 +325,26 @@ def __boa_private_{self.name}__({args_signature}) -> {self.return_type[0]}:
 """
 
 
-class VVMEval(_VVMInternal):
+class VVMInjectedFunction(_VVMInternal):
     """
-    A Vyper eval statement which can be used to evaluate vyper statements
-    via vvm-compiled contracts. This implementation has some drawbacks:
-    - It is very slow, as it requires the complete contract to be recompiled.
-    - It does not detect the return type, as it is currently not possible.
-    - It will temporarily change the bytecode at the contract's address.
+    A Vyper function that is injected into a VVM contract.
+    It will temporarily change the bytecode at the contract's address.
     """
 
-    def __init__(self, code: str, contract: VVMContract, return_type: str = None):
-        abi = {
-            "anonymous": False,
-            "inputs": [],
-            "outputs": ([{"name": "eval", "type": return_type}] if return_type else []),
-            "name": "__boa_debug__",
-            "type": "function",
-        }
-        super().__init__(abi, contract.contract_name)
+    def __init__(self, code: str, contract: VVMContract):
         self.contract = contract
         self.code = code
+        abi = [i for i in self._compiler_output["abi"] if i not in contract.abi]
+        if len(abi) != 1:
+            err = "Expected exactly one new ABI entry after injecting function. "
+            err += f"Found {abi}."
+            raise ValueError(err)
+
+        super().__init__(abi[0], contract.contract_name)
 
     @cached_property
     def source_code(self):
-        debug_body = self.code
-        return_sig = ""
-        if self.return_type:
-            return_sig = f"-> ({', '.join(self.return_type)})"
-            debug_body = f"return {self.code}"
-        return f"""
-@external
-@payable
-def __boa_debug__() {return_sig}:
-    {debug_body}
-"""
+        return self.code
 
 
 def _get_storage_variable_types(spec: dict) -> tuple[list[dict], str]:
diff --git a/tests/unitary/contracts/vvm/test_vvm.py b/tests/unitary/contracts/vvm/test_vvm.py
index 350f4521..ab29ea79 100644
--- a/tests/unitary/contracts/vvm/test_vvm.py
+++ b/tests/unitary/contracts/vvm/test_vvm.py
@@ -1,3 +1,5 @@
+import pytest
+
 import boa
 
 mock_3_10_path = "tests/unitary/contracts/vvm/mock_3_10.vy"
@@ -52,13 +54,35 @@ def test_vvm_internal():
     assert contract._storage.hash_map.get(address, 0) == 69
 
 
-def test_vvm_eval():
+def test_vvm_inject_fn():
     contract = boa.loads(mock_3_10_code, 43)
-    assert contract.eval("self.bar", "uint256") == 43
-    assert contract.eval("self.bar = 44") is None
+    contract.inject_function(
+        """
+@external
+def set_bar(bar: uint256):
+    self.bar = bar
+"""
+    )
+    assert contract.bar() == 43
+    assert contract.set_bar(44) is None
     assert contract.bar() == 44
 
 
+def test_vvm_inject_fn_exists():
+    contract = boa.loads(mock_3_10_code, 43)
+    code = """
+@external
+def bytecode():
+    assert False, "Function injected"
+"""
+    with pytest.raises(ValueError) as e:
+        contract.inject_function(code)
+    assert "Function bytecode already exists" in str(e.value)
+    contract.inject_function(code, force=True)
+    with boa.reverts("Function injected"):
+        contract.bytecode()
+
+
 def test_forward_args_on_deploy():
     with open(mock_3_10_path) as f:
         code = f.read()

From 3f302edad4a1997efd0fadc8fb5de3b6f830b92b Mon Sep 17 00:00:00 2001
From: Charles Cooper <cooper.charles.m@gmail.com>
Date: Sat, 19 Oct 2024 10:48:37 -0400
Subject: [PATCH 16/20] some refactor

---
 boa/contracts/vvm/vvm_contract.py | 45 ++++++++++++++++++++++---------
 boa/interpret.py                  | 10 ++-----
 boa/util/cached_vvm.py            | 26 ------------------
 3 files changed, 34 insertions(+), 47 deletions(-)
 delete mode 100644 boa/util/cached_vvm.py

diff --git a/boa/contracts/vvm/vvm_contract.py b/boa/contracts/vvm/vvm_contract.py
index 6a136dea..195892b9 100644
--- a/boa/contracts/vvm/vvm_contract.py
+++ b/boa/contracts/vvm/vvm_contract.py
@@ -1,18 +1,36 @@
 import re
 from functools import cached_property
 from pathlib import Path
-from typing import Optional
+from typing import Any, Optional
 
+import vvm
 from vyper.utils import method_id
 
 from boa.contracts.abi.abi_contract import ABIContract, ABIContractFactory, ABIFunction
 from boa.environment import Env
 from boa.rpc import to_bytes
-from boa.util import cached_vvm
 from boa.util.abi import Address
+from boa.util.disk_cache import get_disk_cache
 from boa.util.eip5202 import generate_blueprint_bytecode
 
 
+def _compile_source(*args, **kwargs) -> Any:
+    """
+    Compile Vyper source code via the VVM.
+    When a disk cache is available, the result of the compilation is cached.
+    """
+    disk_cache = get_disk_cache()
+
+    def _compile():
+        return vvm.compile_source(*args, **kwargs)
+
+    if disk_cache is None:
+        return _compile()
+
+    cache_key = f"{args}{kwargs}"
+    return disk_cache.caching_lookup(cache_key, _compile)
+
+
 class VVMDeployer(ABIContractFactory):
     """
     A deployer that uses the Vyper Version Manager (VVM).
@@ -46,9 +64,8 @@ def bytecode(self):
         return to_bytes(self.compiler_output["bytecode"])
 
     @classmethod
-    def from_compiler_output(
+    def from_source_code(
         cls,
-        compiler_output: dict,
         source_code: str,
         vyper_version: str,
         filename: Optional[str] = None,
@@ -56,6 +73,9 @@ def from_compiler_output(
     ):
         if name is None:
             name = Path(filename).stem if filename is not None else "<VVMContract>"
+        compiled_src = _compile_source(source_code, vyper_version=vyper_version)
+        compiler_output = compiled_src["<stdin>"]
+
         return cls(name, compiler_output, source_code, vyper_version, filename)
 
     @cached_property
@@ -183,17 +203,18 @@ def internal(self):
         recompiled.
         """
 
-        def internal():
+        # an object with working setattr
+        def _obj():
             return None
 
-        result = cached_vvm.compile_source(
+        result = _compile_source(
             self.source_code, vyper_version=self.vyper_version, output_format="metadata"
         )["function_info"]
         for fn_name, meta in result.items():
             if meta["visibility"] == "internal":
                 function = VVMInternalFunction(meta, self)
-                setattr(internal, function.name, function)
-        return internal
+                setattr(_obj, function.name, function)
+        return _obj
 
 
 class _VVMInternal(ABIFunction):
@@ -212,9 +233,7 @@ def _override_bytecode(self) -> bytes:
     def _compiler_output(self):
         assert isinstance(self.contract, VVMContract)  # help mypy
         source = "\n".join((self.contract.source_code, self.source_code))
-        compiled = cached_vvm.compile_source(
-            source, vyper_version=self.contract.vyper_version
-        )
+        compiled = _compile_source(source, vyper_version=self.contract.vyper_version)
         return compiled["<stdin>"]
 
     @property
@@ -331,9 +350,9 @@ class VVMInjectedFunction(_VVMInternal):
     It will temporarily change the bytecode at the contract's address.
     """
 
-    def __init__(self, code: str, contract: VVMContract):
+    def __init__(self, source_code: str, contract: VVMContract):
         self.contract = contract
-        self.code = code
+        self._source_code = source_code
         abi = [i for i in self._compiler_output["abi"] if i not in contract.abi]
         if len(abi) != 1:
             err = "Expected exactly one new ABI entry after injecting function. "
diff --git a/boa/interpret.py b/boa/interpret.py
index 46faeb2a..d5601bcb 100644
--- a/boa/interpret.py
+++ b/boa/interpret.py
@@ -33,7 +33,6 @@
 from boa.environment import Env
 from boa.explorer import Etherscan, get_etherscan
 from boa.rpc import json
-from boa.util import cached_vvm
 from boa.util.abi import Address
 from boa.util.disk_cache import get_disk_cache
 
@@ -231,8 +230,7 @@ def loads_partial(
 
     version = vvm.detect_vyper_version_from_source(source_code)
     if version is not None and version != Version(vyper.__version__):
-        # TODO: pass name to loads_partial_vvm, not filename
-        return _loads_partial_vvm(source_code, str(version), filename)
+        return _loads_partial_vvm(source_code, str(version), filename, name)
 
     compiler_args = compiler_args or {}
 
@@ -261,11 +259,7 @@ def _loads_partial_vvm(
     # install the requested version if not already installed
     vvm.install_vyper(version=version)
 
-    compiled_src = cached_vvm.compile_source(source_code, vyper_version=version)
-    compiler_output = compiled_src["<stdin>"]
-    return VVMDeployer.from_compiler_output(
-        compiler_output, source_code, version, filename, name
-    )
+    return VVMDeployer.from_source_code(source_code, version, filename, name)
 
 
 def from_etherscan(
diff --git a/boa/util/cached_vvm.py b/boa/util/cached_vvm.py
deleted file mode 100644
index 653543b6..00000000
--- a/boa/util/cached_vvm.py
+++ /dev/null
@@ -1,26 +0,0 @@
-from typing import Any
-
-import vvm
-
-from boa.util.disk_cache import get_disk_cache
-
-
-def compile_source(*args, **kwargs) -> Any:
-    """
-    Compile Vyper source code via the VVM.
-    When a disk cache is available, the result of the compilation is cached.
-    Note the cache only works if the function is called the same way (args/kwargs).
-    :param args: Arguments to pass to vvm.compile_source
-    :param kwargs: Keyword arguments to pass to vvm.compile_source
-    :return: Compilation output
-    """
-    disk_cache = get_disk_cache()
-
-    def _compile():
-        return vvm.compile_source(*args, **kwargs)
-
-    if disk_cache is None:
-        return _compile()
-
-    cache_key = f"{args}{kwargs}"
-    return disk_cache.caching_lookup(cache_key, _compile)

From da7467ef83b568940efb97836d77e1f4f26ffcfe Mon Sep 17 00:00:00 2001
From: Charles Cooper <cooper.charles.m@gmail.com>
Date: Sat, 19 Oct 2024 10:55:12 -0400
Subject: [PATCH 17/20] fix bytecode override

---
 boa/contracts/abi/abi_contract.py |  5 +++++
 boa/contracts/vvm/vvm_contract.py | 11 ++---------
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/boa/contracts/abi/abi_contract.py b/boa/contracts/abi/abi_contract.py
index 47f8b382..2579f68f 100644
--- a/boa/contracts/abi/abi_contract.py
+++ b/boa/contracts/abi/abi_contract.py
@@ -124,6 +124,10 @@ def __call__(self, *args, value=0, gas=None, sender=None, **kwargs):
         if not self.contract or not self.contract.env:
             raise Exception(f"Cannot call {self} without deploying contract.")
 
+        override_bytecode = None
+        if hasattr(self, "_override_bytecode"):
+            override_bytecode = self._override_bytecode
+
         computation = self.contract.env.execute_code(
             to_address=self.contract.address,
             sender=sender,
@@ -132,6 +136,7 @@ def __call__(self, *args, value=0, gas=None, sender=None, **kwargs):
             gas=gas,
             is_modifying=self.is_mutable,
             contract=self.contract,
+            override_bytecode=override_bytecode,
         )
 
         match self.contract.marshal_to_python(computation, self.return_type):
diff --git a/boa/contracts/vvm/vvm_contract.py b/boa/contracts/vvm/vvm_contract.py
index 195892b9..b6a2ba8a 100644
--- a/boa/contracts/vvm/vvm_contract.py
+++ b/boa/contracts/vvm/vvm_contract.py
@@ -244,15 +244,6 @@ def source_code(self) -> str:
         """
         raise NotImplementedError
 
-    def __call__(self, *args, **kwargs):
-        env = self.contract.env
-        assert isinstance(self.contract, VVMContract)  # help mypy
-        env.set_code(self.contract.address, self._override_bytecode)
-        try:
-            return super().__call__(*args, **kwargs)
-        finally:
-            env.set_code(self.contract.address, self.contract.bytecode_runtime)
-
 
 class VVMInternalFunction(_VVMInternal):
     """
@@ -324,6 +315,8 @@ def __init__(self, name, spec, contract):
         self.contract = contract
 
     def get(self, *args):
+        # get the value of the storage variable. note that this is
+        # different from the behavior of VyperContract storage variables!
         return self.__call__(*args)
 
     @cached_property

From d6dad6387a9c3998051e375480de78502f923ad7 Mon Sep 17 00:00:00 2001
From: Charles Cooper <cooper.charles.m@gmail.com>
Date: Sat, 19 Oct 2024 11:00:42 -0400
Subject: [PATCH 18/20] fix API regression

---
 boa/interpret.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/boa/interpret.py b/boa/interpret.py
index d5601bcb..a768c67e 100644
--- a/boa/interpret.py
+++ b/boa/interpret.py
@@ -34,7 +34,8 @@
 from boa.explorer import Etherscan, get_etherscan
 from boa.rpc import json
 from boa.util.abi import Address
-from boa.util.disk_cache import get_disk_cache
+# export set_cache_dir, NOTE: consider moving to boa/__init__.py
+from boa.util.disk_cache import get_disk_cache, set_cache_dir  # noqa: F401
 
 if TYPE_CHECKING:
     from vyper.semantics.analysis.base import ImportInfo

From 550b33415c03a36b3e2bfee6300b19cd458360f0 Mon Sep 17 00:00:00 2001
From: Charles Cooper <cooper.charles.m@gmail.com>
Date: Mon, 21 Oct 2024 11:10:41 -0400
Subject: [PATCH 19/20] fix lint

---
 boa/interpret.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/boa/interpret.py b/boa/interpret.py
index a768c67e..cf014b88 100644
--- a/boa/interpret.py
+++ b/boa/interpret.py
@@ -34,6 +34,7 @@
 from boa.explorer import Etherscan, get_etherscan
 from boa.rpc import json
 from boa.util.abi import Address
+
 # export set_cache_dir, NOTE: consider moving to boa/__init__.py
 from boa.util.disk_cache import get_disk_cache, set_cache_dir  # noqa: F401
 

From ea4bccb119c9d2f6051b0f85cea903578c6fff61 Mon Sep 17 00:00:00 2001
From: Charles Cooper <cooper.charles.m@gmail.com>
Date: Mon, 21 Oct 2024 11:12:42 -0400
Subject: [PATCH 20/20] fix bad variable

---
 boa/contracts/vvm/vvm_contract.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/boa/contracts/vvm/vvm_contract.py b/boa/contracts/vvm/vvm_contract.py
index b6a2ba8a..377bb8d8 100644
--- a/boa/contracts/vvm/vvm_contract.py
+++ b/boa/contracts/vvm/vvm_contract.py
@@ -356,7 +356,7 @@ def __init__(self, source_code: str, contract: VVMContract):
 
     @cached_property
     def source_code(self):
-        return self.code
+        return self._source_code
 
 
 def _get_storage_variable_types(spec: dict) -> tuple[list[dict], str]: