From c44e103cf582b2b0eaa6dc28b836c85e624bb538 Mon Sep 17 00:00:00 2001
From: Alberto <albicento.ac@gmail.com>
Date: Fri, 6 Sep 2024 14:12:33 +0200
Subject: [PATCH] feat: cache vvm contracts (#295)

* feat: cache vvm contracts

---------

Co-authored-by: Charles Cooper <cooper.charles.m@gmail.com>
---
 boa/interpret.py                  | 21 ++++++++++---
 tests/unitary/utils/test_cache.py | 51 ++++++++++++++++++++++++++++++-
 2 files changed, 66 insertions(+), 6 deletions(-)

diff --git a/boa/interpret.py b/boa/interpret.py
index 8386d1fb..db746f9f 100644
--- a/boa/interpret.py
+++ b/boa/interpret.py
@@ -233,13 +233,24 @@ def load_partial(filename: str, compiler_args=None):
 
 
 def _loads_partial_vvm(source_code: str, version: str, filename: str):
-    # will install the request version if not already installed
+    global _disk_cache
+
+    # install the requested 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)
+    def _compile():
+        compiled_src = vvm.compile_source(source_code, vyper_version=version)
+        compiler_output = compiled_src["<stdin>"]
+        return VVMDeployer.from_compiler_output(compiler_output, filename=filename)
+
+    # 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)
 
 
 def from_etherscan(
diff --git a/tests/unitary/utils/test_cache.py b/tests/unitary/utils/test_cache.py
index 27a731f3..64659821 100644
--- a/tests/unitary/utils/test_cache.py
+++ b/tests/unitary/utils/test_cache.py
@@ -1,7 +1,20 @@
+from unittest.mock import patch
+
+import pytest
 from vyper.compiler import CompilerData
 
 from boa.contracts.vyper.vyper_contract import VyperDeployer
-from boa.interpret import _disk_cache, compiler_data
+from boa.interpret import _disk_cache, _loads_partial_vvm, compiler_data, set_cache_dir
+
+
+@pytest.fixture(autouse=True)
+def cache_dir(tmp_path):
+    tmp = _disk_cache.cache_dir
+    try:
+        set_cache_dir(tmp_path)
+        yield
+    finally:
+        set_cache_dir(tmp)
 
 
 def test_cache_contract_name():
@@ -17,6 +30,42 @@ def test_cache_contract_name():
     assert str(test2.contract_path) == "test2"
 
 
+def test_cache_vvm():
+    code = """
+x: constant(int128) = 1000
+"""
+    version = "0.2.8"
+    version2 = "0.3.1"
+    assert _disk_cache is not None
+
+    # Mock vvm.compile_source
+    with patch("vvm.compile_source") as mock_compile:
+        # Set up the mock to return a valid compiler output
+        mock_compile.return_value = {"<stdin>": {"abi": [], "bytecode": "0x1234"}}
+
+        assert mock_compile.call_count == 0
+
+        # First call should hit vvm.compile_source
+        test1 = _loads_partial_vvm(code, version, "fake_file.vy")
+        assert mock_compile.call_count == 1
+
+        # Second call should hit the cache
+        test2 = _loads_partial_vvm(code, version, "fake_file.vy")
+        assert mock_compile.call_count == 1
+
+        # using a different filename should also hit the cache
+        test3 = _loads_partial_vvm(code, version, "fake_fileeeee.vy")
+        assert mock_compile.call_count == 1
+
+        # using a different vyper version should *miss* the cache
+        _loads_partial_vvm(code, version2, "fake_file.vy")
+        assert mock_compile.call_count == 2
+
+    assert test1.abi == test2.abi == test3.abi
+    assert test1.bytecode == test2.bytecode == test3.bytecode
+    assert test1.filename == test2.filename
+
+
 def _to_dict(data: CompilerData) -> dict:
     """
     Serialize the `CompilerData` object to a dictionary for comparison.