Implement Colab support

boa_zksync/__init__.py

@@ -30,3 +30,4 @@ 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

@@ -1,5 +1,12 @@
+import logging
+import os
+from shutil import which
+
+import requests
+from boa.rpc import EthereumRPC
+
 try:
-    from boa.integrations.jupyter import BrowserEnv
+    from boa.integrations.jupyter.browser import BrowserEnv, BrowserSigner, BrowserRPC, colab_eval_js
 except ImportError:
     raise ModuleNotFoundError(
         "The `BrowserEnv` class requires Jupyter to be installed. "
@@ -9,7 +16,71 @@
 from boa_zksync.environment import ZksyncEnv
 
 
-class ZksyncBrowserEnv(ZksyncEnv, BrowserEnv):
+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.")
+            install_zkvyper_compiler()
+
+        super().__init__(rpc=BrowserRPC(), *args, **kwargs)
+        self.signer = BrowserSigner(address)
+        self.set_eoa(self.signer)
+
+    def set_chain_id(self, chain_id: int | str):
+        self._rpc.fetch(
+            "wallet_switchEthereumChain",
+            [{"chainId": chain_id if isinstance(chain_id, str) else hex(chain_id)}],
+        )
+        self._reset_fork()
+
+    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.")
+            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",
+    destination="/usr/local/bin/zkvyper",
+):
+    """
+    Downloads the zkvyper binary from the given source URL and installs it to
+    the destination directory.
+
+    This is a very basic implementation - usually users want to install the binary
+    manually, but in the Colab environment, we can automate this process.
+    """
+    response = requests.get(source)
+    with open(destination, "wb") as f:
+        f.write(response.content)
+
+    os.chmod(destination, 0o755)  # make it executable
+    assert os.system("zkvyper --version") == 0  # check if it works
+
+
+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",
+    destination="/usr/local/bin/era_test_node",
+):
+    """
+    Downloads the era-test-node binary from the given source URL and installs it to
+    the destination directory.
+
+    This is a very basic implementation - usually users want to install the binary
+    manually, but in the Colab environment, we can automate this process.
+    """
+    response = requests.get(source)
+    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(f"mv era_test_node {destination}")
+    os.system(f"{destination} --version")
+    os.system("rm era_test_node.tar.gz")

boa_zksync/deployer.py

@@ -23,17 +23,19 @@ def __init__(self, compiler_data: ZksyncCompilerData, name: str, filename: str):
         self.compiler_data = compiler_data
 
     def deploy(self, *args, value=0, **kwargs):
-        env = Env.get_singleton()
-
         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()
         address, _ = env.deploy_code(
-            bytecode=initcode, value=value, constructor_calldata=constructor_calldata
+            bytecode=bytecode, value=value, constructor_calldata=constructor_calldata
         )
         return ABIContract(
             self._name,

boa_zksync/environment.py

@@ -1,23 +1,21 @@
-import sys
+from collections import namedtuple
 from contextlib import contextmanager
 from functools import cached_property
 from hashlib import sha256
 from pathlib import Path
-from subprocess import Popen
-from typing import Any, Callable, Optional, cast
+from typing import Any, Optional, Iterable
 
 from boa.contracts.abi.abi_contract import ABIContract, ABIContractFactory
 from boa.environment import _AddressType
 from boa.interpret import json
-from boa.network import NetworkEnv, TransactionSettings, _EstimateGasFailed
+from boa.network import NetworkEnv, _EstimateGasFailed
 from boa.rpc import RPC, fixup_dict, to_bytes, to_hex, EthereumRPC
 from boa.util.abi import Address
 from eth.constants import ZERO_ADDRESS
 from eth.exceptions import VMError
-from eth_account import Account
 
 from boa_zksync.node import EraTestNode
-from boa_zksync.util import DeployTransaction, find_free_port, wait_url, stop_subprocess
+from boa_zksync.util import DeployTransaction
 
 _CONTRACT_DEPLOYER_ADDRESS = "0x0000000000000000000000000000000000008006"
 with open(Path(__file__).parent / "IContractDeployer.json") as f:
@@ -46,11 +44,9 @@ def create(self):
 
     def _reset_fork(self, block_identifier="latest"):
         if isinstance(self._rpc, EraTestNode):
-            url = self._rpc._inner_url
+            rpc = self._rpc.inner_rpc
             del self._rpc
-        else:
-            url = self._rpc._rpc_url
-        self._rpc = EthereumRPC(url)
+            self._rpc = rpc
 
     def fork_rpc(self, rpc: EthereumRPC, reset_traces=True, block_identifier="safe", **kwargs):
         """
@@ -100,8 +96,7 @@ def execute_code(
         :param contract: The contract ABI.
         :return: The return value of the contract function.
         """
-        sender = self._check_sender(self._get_sender(sender))
-
+        sender = str(self._check_sender(self._get_sender(sender)))
         hexdata = to_hex(data)
 
         if not is_modifying:
@@ -133,6 +128,7 @@ def deploy_code(
         value=0,
         bytecode=b"",
         constructor_calldata=b"",
+        dependency_bytecodes: Iterable[bytes] = (),
         salt=b"\0" * 32,
         **kwargs,
     ) -> tuple[Address, bytes]:
@@ -143,6 +139,7 @@ def deploy_code(
         :param value: The amount of value to send with the transaction.
         :param bytecode: The bytecode of the contract to deploy.
         :param constructor_calldata: The calldata for the contract constructor.
+        :param dependency_bytecodes: The bytecodes of the blueprints.
         :param salt: The salt for the contract deployment.
         :param kwargs: Additional parameters for the transaction.
         :return: The address of the deployed contract and the bytecode hash.
@@ -171,6 +168,8 @@ def deploy_code(
             ),
             bytecode=bytecode,
             bytecode_hash=bytecode_hash,
+            dependency_bytecodes=list(dependency_bytecodes),
+            dependency_bytecode_hashes=[_hash_code(bc) for bc in dependency_bytecodes],
             chain_id=chain_id,
             paymaster_params=kwargs.pop("paymaster_params", None),
         )

boa_zksync/node.py

@@ -9,7 +9,7 @@
 
 class EraTestNode(EthereumRPC):
     def __init__(self, rpc: EthereumRPC, block_identifier="safe"):
-        self._inner_url = rpc._rpc_url
+        self.inner_rpc = rpc
 
         port = find_free_port()
         fork_at = [
@@ -21,16 +21,12 @@ def __init__(self, rpc: EthereumRPC, block_identifier="safe"):
             "--port",
             f"{port}",
             "fork",
-            self._inner_url,
+            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}")
         wait_url(self._rpc_url)
 
-    @property
-    def inner_url(self):
-        return self._inner_url
-
     def __del__(self):
         stop_subprocess(self._test_node)

boa_zksync/util.py

@@ -57,10 +57,13 @@ class DeployTransaction:
     calldata: bytes  # the result of calling the deployer's create.prepare_calldata
     bytecode: bytes
     bytecode_hash: bytes
+    dependency_bytecodes: list[bytes]
+    dependency_bytecode_hashes: list[bytes]
     chain_id: int
     paymaster_params: tuple[int, bytes] | None
 
     def get_estimate_tx(self):
+        bytecodes = [self.bytecode] + self.dependency_bytecodes
         return {
             "transactionType": f"0x{_EIP712_TYPE.hex()}",
             "chain_id": self.chain_id,
@@ -74,7 +77,7 @@ def get_estimate_tx(self):
             "data": f"0x{self.calldata.hex()}",
             "eip712Meta": {
                 "gasPerPubdata": f"0x{_GAS_PER_PUB_DATA_DEFAULT:0x}",
-                "factoryDeps": [[int(b) for b in self.bytecode]],
+                "factoryDeps": [[int(byte) for byte in bytecode] for bytecode in bytecodes],
             },
         }
 
@@ -100,7 +103,7 @@ def sign_typed_data(
                 "nonce": self.nonce,
                 "value": self.value,
                 "data": self.calldata,
-                "factoryDeps": [self.bytecode_hash],
+                "factoryDeps": [self.bytecode_hash] + self.dependency_bytecode_hashes,
                 "paymaster": paymaster,
                 "paymasterInput": paymaster_input,
             },
@@ -120,6 +123,7 @@ def rlp_encode(self, signature: SignedMessage | str, estimated_gas: int) -> byte
             elements=([_BINARY, _BINARY] if self.paymaster_params else None),
             strict=False,
         )
+        bytecodes = [self.bytecode] + self.dependency_bytecodes
         return _EIP712_TYPE + rlp.encode(
             [
                 _INT.serialize(self.nonce),
@@ -135,7 +139,7 @@ def rlp_encode(self, signature: SignedMessage | str, estimated_gas: int) -> byte
                 _INT.serialize(self.chain_id),
                 _BINARY.serialize(to_bytes(self.sender)),
                 _INT.serialize(_GAS_PER_PUB_DATA_DEFAULT),
-                List(elements=([_BINARY]), strict=False).serialize([self.bytecode]),
+                List(elements=([_BINARY]), strict=False).serialize(bytecodes),
                 _BINARY.serialize(
                     to_bytes(signature)
                     if isinstance(signature, str)

tests/test_deploy.py

@@ -1,31 +1,11 @@
-import os
 import sys
 from subprocess import Popen
 
 import pytest
-from eth_account import Account
 
 from boa_zksync.interpret import loads_zksync
 from boa_zksync.util import find_free_port, wait_url, stop_subprocess
 
-code = """
-totalSupply: public(uint256)
-balances: HashMap[address, uint256]
-
-@external
-def __init__(t: uint256):
-    self.totalSupply = t
-    self.balances[self] = t
-
-@external
-def update_total_supply(t: uint16):
-    self.totalSupply += convert(t, uint256)
-
-@external
-def raise_exception(t: uint256):
-    raise "oh no!"
-"""
-
 STARTING_SUPPLY = 100
 
 
@@ -44,6 +24,23 @@ def era_test_node():
 
 @pytest.fixture(scope="module")
 def simple_contract():
+    code = """
+totalSupply: public(uint256)
+balances: HashMap[address, uint256]
+
+@external
+def __init__(t: uint256):
+    self.totalSupply = t
+    self.balances[self] = t
+
+@external
+def update_total_supply(t: uint16):
+    self.totalSupply += convert(t, uint256)
+
+@external
+def raise_exception(t: uint256):
+    raise "oh no!"
+"""
     return loads_zksync(code, STARTING_SUPPLY)