From 2badb067c9bf1a10ebfe3b2decf75add6daa8e43 Mon Sep 17 00:00:00 2001
From: Charles Cooper <cooper.charles.m@gmail.com>
Date: Mon, 21 Oct 2024 11:35:29 -0400
Subject: [PATCH] fix stomp() (#337)

this was a regression in 287523f7c176d14ade84c9053aacd28a33555289

add a regression test
---
 boa/contracts/vyper/vyper_contract.py         |  6 +--
 .../contracts/vyper/test_vyper_contract.py    | 51 +++++++++++++++++++
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/boa/contracts/vyper/vyper_contract.py b/boa/contracts/vyper/vyper_contract.py
index cdf409c8..d7027a25 100644
--- a/boa/contracts/vyper/vyper_contract.py
+++ b/boa/contracts/vyper/vyper_contract.py
@@ -100,8 +100,8 @@ def stomp(self, address: Any, data_section=None) -> "VyperContract":
         address = Address(address)
 
         ret = self.deploy(override_address=address, skip_initcode=True)
-        vm = ret.env.vm
-        old_bytecode = vm.state.get_code(address.canonical_address)
+        vm = ret.env.evm
+        old_bytecode = vm.get_code(address)
         new_bytecode = self.compiler_data.bytecode_runtime
 
         immutables_size = self.compiler_data.global_ctx.immutable_section_bytes
@@ -109,7 +109,7 @@ def stomp(self, address: Any, data_section=None) -> "VyperContract":
             data_section = old_bytecode[-immutables_size:]
             new_bytecode += data_section
 
-        vm.state.set_code(address.canonical_address, new_bytecode)
+        vm.set_code(address, new_bytecode)
         ret.env.register_contract(address, ret)
         ret._set_bytecode(new_bytecode)
         return ret
diff --git a/tests/unitary/contracts/vyper/test_vyper_contract.py b/tests/unitary/contracts/vyper/test_vyper_contract.py
index b522e935..70370082 100644
--- a/tests/unitary/contracts/vyper/test_vyper_contract.py
+++ b/tests/unitary/contracts/vyper/test_vyper_contract.py
@@ -72,3 +72,54 @@ def foo() -> bool:
     c = boa.loads(code)
 
     c.foo()
+
+
+def test_stomp():
+    code1 = """
+VAR: immutable(uint256)
+
+@deploy
+def __init__():
+    VAR = 12345
+
+@external
+def foo() -> uint256:
+    return VAR
+
+@external
+def bar() -> bool:
+    return True
+    """
+    code2 = """
+VAR: immutable(uint256)
+
+@deploy
+def __init__():
+    VAR = 12345
+
+@external
+def foo() -> uint256:
+    return VAR
+
+@external
+def bar() -> bool:
+    return False
+    """
+
+    deployer = boa.loads_partial(code1)
+
+    c = deployer.deploy()
+
+    assert c.foo() == 12345
+    assert c.bar() is True
+
+    deployer2 = boa.loads_partial(code2)
+
+    c2 = deployer2.stomp(c.address)
+
+    assert c2.foo() == 12345
+    assert c2.bar() is False
+
+    # the bytecode at the original contract has been stomped :scream:
+    assert c.foo() == 12345
+    assert c.bar() is False