Implement blueprints

Note: trace.returndata is incorrect, we are researching what is the right way.

Makefile

@@ -10,6 +10,7 @@ mypy:
 		--follow-imports=silent \
 		--ignore-missing-imports \
 		--implicit-optional \
+		--install-types \
 		-p boa_zksync
 
 black:

boa_zksync/__init__.py

@@ -30,4 +30,3 @@ def set_zksync_browser_env(address=None):
 boa.eval_zksync = eval_zksync
 boa.set_zksync_env = set_zksync_env
 boa.set_zksync_browser_env = set_zksync_browser_env
-

boa_zksync/browser.py

@@ -3,29 +3,25 @@
 from shutil import which
 
 import requests
+from boa.integrations.jupyter.browser import BrowserRPC, BrowserSigner, colab_eval_js
 from boa.rpc import EthereumRPC
 
-try:
-    from boa.integrations.jupyter.browser import BrowserEnv, BrowserSigner, BrowserRPC, colab_eval_js
-except ImportError:
-    raise ModuleNotFoundError(
-        "The `BrowserEnv` class requires Jupyter to be installed. "
-        "Please be careful when importing the browser files outside of Jupyter env."
-    )
-
 from boa_zksync.environment import ZksyncEnv
 
 
 class ZksyncBrowserEnv(ZksyncEnv):
     """
     A zkSync environment for deploying contracts using a browser wallet RPC.
     """
+
     def __init__(self, address=None, *args, **kwargs):
         if colab_eval_js and not which("zkvyper"):
-            logging.warning("Automatically installing zkvyper compiler in the Colab environment.")
+            logging.warning(
+                "Automatically installing zkvyper compiler in the Colab environment."
+            )
             install_zkvyper_compiler()
 
-        super().__init__(rpc=BrowserRPC(), *args, **kwargs)
+        super().__init__(BrowserRPC(), *args, **kwargs)
         self.signer = BrowserSigner(address)
         self.set_eoa(self.signer)
 
@@ -36,17 +32,21 @@ def set_chain_id(self, chain_id: int | str):
         )
         self._reset_fork()
 
-    def fork_rpc(self, rpc: EthereumRPC, reset_traces=True, block_identifier="safe", **kwargs):
+    def fork_rpc(
+        self, rpc: EthereumRPC, reset_traces=True, block_identifier="safe", **kwargs
+    ):
         if colab_eval_js and not which("era_test_node"):
-            logging.warning("Automatically installing era-test-node in the Colab environment.")
+            logging.warning(
+                "Automatically installing era-test-node in the Colab environment."
+            )
             install_era_test_node()
 
         return super().fork_rpc(rpc, reset_traces, block_identifier, **kwargs)
 
 
 def install_zkvyper_compiler(
     source="https://raw.githubusercontent.com/matter-labs/zkvyper-bin/"
-           "66cc159d9b6af3b5616f6ed7199bd817bf42bf0a/linux-amd64/zkvyper-linux-amd64-musl-v1.4.0",
+    "66cc159d9b6af3b5616f6ed7199bd817bf42bf0a/linux-amd64/zkvyper-linux-amd64-musl-v1.4.0",
     destination="/usr/local/bin/zkvyper",
 ):
     """
@@ -66,7 +66,7 @@ def install_zkvyper_compiler(
 
 def install_era_test_node(
     source="https://github.com/matter-labs/era-test-node/releases/download/"
-           "v0.1.0-alpha.19/era_test_node-v0.1.0-alpha.19-x86_64-unknown-linux-gnu.tar.gz",
+    "v0.1.0-alpha.19/era_test_node-v0.1.0-alpha.19-x86_64-unknown-linux-gnu.tar.gz",
     destination="/usr/local/bin/era_test_node",
 ):
     """
@@ -80,7 +80,7 @@ def install_era_test_node(
     with open("era_test_node.tar.gz", "wb") as f:
         f.write(response.content)
 
-    os.system(f"tar --extract --file=era_test_node.tar.gz")
+    os.system("tar --extract --file=era_test_node.tar.gz")
     os.system(f"mv era_test_node {destination}")
     os.system(f"{destination} --version")
     os.system("rm era_test_node.tar.gz")

boa_zksync/deployer.py

@@ -6,6 +6,7 @@
 from boa.util.abi import Address
 
 from boa_zksync.compile import ZksyncCompilerData
+from boa_zksync.environment import ZksyncEnv
 
 
 class ZksyncDeployer(ABIContractFactory):
@@ -23,28 +24,38 @@ def __init__(self, compiler_data: ZksyncCompilerData, name: str, filename: str):
         self.compiler_data = compiler_data
 
     def deploy(self, *args, value=0, **kwargs):
-        initcode = to_bytes(self.compiler_data.bytecode)
-        return self._deploy(initcode, *args, value=value, **kwargs)
-
-    def _deploy(self, bytecode, *args, value=0, dependency_bytecodes=(), **kwargs):
-        constructor_calldata = (
-            self.constructor.prepare_calldata(*args, **kwargs)
-            if args or kwargs
-            else b""
-        )
-
         env = Env.get_singleton()
+        assert isinstance(
+            env, ZksyncEnv
+        ), "ZksyncDeployer can only be used in zkSync environments"
+
         address, _ = env.deploy_code(
-            bytecode=bytecode, value=value, constructor_calldata=constructor_calldata
+            bytecode=to_bytes(self.compiler_data.bytecode),
+            value=value,
+            constructor_calldata=(
+                self.constructor.prepare_calldata(*args, **kwargs)
+                if args or kwargs
+                else b""
+            ),
         )
-        return ABIContract(
+        address = Address(address)
+        abi_contract = ABIContract(
             self._name,
             self.abi,
             self._functions,
-            address=Address(address),
+            address=address,
             filename=self._filename,
             env=env,
         )
+        env.register_contract(address, self)
+        return abi_contract
+
+    def deploy_as_blueprint(self, *args, **kwargs):
+        """
+        In zkSync, any contract can be used as a blueprint.
+        Note that we do need constructor arguments for this.
+        """
+        return self.deploy(*args, **kwargs)
 
     @cached_property
     def constructor(self):

boa_zksync/environment.py

@@ -1,15 +1,16 @@
 from collections import namedtuple
 from contextlib import contextmanager
+from dataclasses import dataclass, field
 from functools import cached_property
 from hashlib import sha256
 from pathlib import Path
-from typing import Any, Optional, Iterable
+from typing import Any, Iterable, Optional
 
 from boa.contracts.abi.abi_contract import ABIContract, ABIContractFactory
 from boa.environment import _AddressType
 from boa.interpret import json
 from boa.network import NetworkEnv, _EstimateGasFailed
-from boa.rpc import RPC, fixup_dict, to_bytes, to_hex, EthereumRPC
+from boa.rpc import RPC, EthereumRPC, fixup_dict, to_bytes, to_hex
 from boa.util.abi import Address
 from eth.constants import ZERO_ADDRESS
 from eth.exceptions import VMError
@@ -48,7 +49,9 @@ def _reset_fork(self, block_identifier="latest"):
             del self._rpc
             self._rpc = rpc
 
-    def fork_rpc(self, rpc: EthereumRPC, reset_traces=True, block_identifier="safe", **kwargs):
+    def fork_rpc(
+        self, rpc: EthereumRPC, reset_traces=True, block_identifier="safe", **kwargs
+    ):
         """
         Fork the environment to a local chain.
         :param rpc: RPC to fork from
@@ -58,8 +61,8 @@ def fork_rpc(self, rpc: EthereumRPC, reset_traces=True, block_identifier="safe",
         """
         self._reset_fork(block_identifier)
         if reset_traces:
-            self.sha3_trace = {}
-            self.sstore_trace = {}
+            self.sha3_trace: dict = {}
+            self.sstore_trace: dict = {}
         self._rpc = EraTestNode(rpc, block_identifier)
 
     def register_contract(self, address, obj):
@@ -99,27 +102,34 @@ def execute_code(
         sender = str(self._check_sender(self._get_sender(sender)))
         hexdata = to_hex(data)
 
+        args = {
+            "from": sender,
+            "to": to_address,
+            "gas": gas,
+            "value": value,
+            "data": hexdata,
+        }
         if not is_modifying:
-            args = fixup_dict(
-                {
-                    "from": sender,
-                    "to": to_address,
-                    "gas": gas,
-                    "value": value,
-                    "data": hexdata,
-                }
-            )
-            output = self._rpc.fetch("eth_call", [args, "latest"])
-            return ZksyncComputation(to_bytes(output))
+            output = self._rpc.fetch("eth_call", [fixup_dict(args), "latest"])
+            return ZksyncComputation(args, to_bytes(output))
 
         try:
             receipt, trace = self._send_txn(
                 from_=sender, to=to_address, value=value, gas=gas, data=hexdata
             )
         except _EstimateGasFailed:
-            return ZksyncComputation(error=VMError("Estimate gas failed"))
+            return ZksyncComputation(args, error=VMError("Estimate gas failed"))
+
+        try:
+            # when calling create_from_blueprint, the address is not returned
+            # we get it from the logs by searching for the event. todo: remove this hack
+            deploy_topic = '0x290afdae231a3fc0bbae8b1af63698b0a1d79b21ad17df0342dfb952fe74f8e5'
+            output = next(x['topics'][3] for x in receipt['logs'] if x['topics'][0] == deploy_topic)
+        except StopIteration:
+            # TODO: This does not return the correct value either.
+            output = trace.returndata
 
-        return ZksyncComputation(receipt.get("output", b""))
+        return ZksyncComputation(args, to_bytes(output))
 
     def deploy_code(
         self,
@@ -174,11 +184,12 @@ def deploy_code(
             paymaster_params=kwargs.pop("paymaster_params", None),
         )
 
-        estimated_gas = int(
-            self._rpc.fetch("eth_estimateGas", [tx.get_estimate_tx()]), 16
-        )
+        estimated_gas = self._rpc.fetch("eth_estimateGas", [tx.get_estimate_tx()])
+        estimated_gas = int(estimated_gas, 16)
+
         signature = tx.sign_typed_data(self._accounts[sender], estimated_gas)
         raw_tx = tx.rlp_encode(signature, estimated_gas)
+
         tx_hash = self._rpc.fetch("eth_sendRawTransaction", ["0x" + raw_tx.hex()])
         receipt = self._rpc.wait_for_tx_receipt(tx_hash, self.tx_settings.poll_timeout)
         return Address(receipt["contractAddress"]), bytecode
@@ -200,34 +211,26 @@ def _hash_code(bytecode: bytes) -> bytes:
     return b"\x01\00" + bytecode_size.to_bytes(2, byteorder="big") + bytecode_hash[4:]
 
 
+@dataclass
 class ZksyncComputation:
-    def __init__(self, output: bytes | None = None, error: VMError | None = None):
-        self._output = output
-        self._error = error
+    args: dict
+    output: bytes | None = None
+    error: VMError | None = None
+    children: list["ZksyncComputation"] = field(default_factory=list)
 
     @property
     def is_success(self) -> bool:
         """
         Return ``True`` if the computation did not result in an error.
         """
-        return self._error is None
+        return self.error is None
 
     @property
     def is_error(self) -> bool:
         """
         Return ``True`` if the computation resulted in an error.
         """
-        return self._error is not None
-
-    @property
-    def error(self) -> VMError:
-        """
-        Return the :class:`~eth.exceptions.VMError` of the computation.
-        Raise ``AttributeError`` if no error exists.
-        """
-        if self._error is None:
-            raise AttributeError("No error exists for this computation")
-        return self._error
+        return self.error is not None
 
     def raise_if_error(self) -> None:
         """
@@ -239,10 +242,7 @@ def raise_if_error(self) -> None:
             raise self.error
 
     @property
-    def output(self) -> bytes:
-        """
-        Get the return value of the computation.
-        """
-        if self._output is None:
-            raise AttributeError("No output exists for this computation")
-        return self._output
+    def msg(self):
+        Message = namedtuple('Message', ['sender','to','gas','value','data'])
+        args = self.args.copy()
+        return Message(args.pop("from"), data=to_bytes(args.pop("data")), **args)

boa_zksync/node.py

@@ -4,25 +4,23 @@
 
 from boa.rpc import EthereumRPC
 
-from boa_zksync.util import stop_subprocess, find_free_port, wait_url
+from boa_zksync.util import find_free_port, stop_subprocess, wait_url
 
 
 class EraTestNode(EthereumRPC):
     def __init__(self, rpc: EthereumRPC, block_identifier="safe"):
         self.inner_rpc = rpc
 
         port = find_free_port()
-        fork_at = [
-            "--fork-at",
-            block_identifier
-        ] if isinstance(block_identifier, int) else []
-        self._test_node = Popen([
-            "era_test_node",
-            "--port",
-            f"{port}",
-            "fork",
-            self.inner_rpc._rpc_url,
-        ] + fork_at, stdout=sys.stdout, stderr=sys.stderr)
+        fork_at = (
+            ["--fork-at", block_identifier] if isinstance(block_identifier, int) else []
+        )
+        self._test_node = Popen(
+            ["era_test_node", "--port", f"{port}", "fork", self.inner_rpc._rpc_url]
+            + fork_at,
+            stdout=sys.stdout,
+            stderr=sys.stderr,
+        )
 
         super().__init__(f"http://localhost:{port}")
         logging.info(f"Started fork node at {self._rpc_url}")

boa_zksync/util.py

@@ -10,8 +10,8 @@
 from eth_account import Account
 from eth_account.datastructures import SignedMessage
 from eth_account.messages import encode_typed_data
-from rlp.sedes import BigEndianInt, Binary, List
 from requests.exceptions import ConnectionError
+from rlp.sedes import BigEndianInt, Binary, List
 
 _EIP712_TYPE = bytes.fromhex("71")
 _EIP712_TYPES_SPEC = {
@@ -77,7 +77,9 @@ def get_estimate_tx(self):
             "data": f"0x{self.calldata.hex()}",
             "eip712Meta": {
                 "gasPerPubdata": f"0x{_GAS_PER_PUB_DATA_DEFAULT:0x}",
-                "factoryDeps": [[int(byte) for byte in bytecode] for bytecode in bytecodes],
+                "factoryDeps": [
+                    [int(byte) for byte in bytecode] for bytecode in bytecodes
+                ],
             },
         }
 

tests/conftest.py

@@ -16,6 +16,12 @@ def zksync_env(rpc, account):
     env = ZksyncEnv(rpc)
     env.add_account(account)
     with boa.swap_env(env):
+        yield env
+
+
+@pytest.fixture(autouse=True)
+def cleanup_env(zksync_env):
+    with zksync_env.anchor():
         yield