Implement internal functions

Makefile

@@ -11,6 +11,7 @@ mypy:
 		--ignore-missing-imports \
 		--implicit-optional \
 		--install-types \
+		--non-interactive \
 		-p boa_zksync
 
 black:

README.md

@@ -92,3 +92,6 @@ deployer.at(address) # Connect a contract to an existing address
 # Run the given source code directly
 boa.eval_zksync("source code")
 ```
+
+### Limitations
+- `# pragma optimize gas` is not supported by Zksync

boa_zksync/__init__.py

@@ -1,12 +1,17 @@
 import boa
 
 from boa_zksync.environment import ZksyncEnv
+from boa_zksync.node import EraTestNode
 
 
 def set_zksync_env(url):
     boa.set_env(ZksyncEnv.from_url(url))
 
 
+def set_zksync_test_env():
+    boa.set_env(ZksyncEnv(rpc=EraTestNode()))
+
+
 def set_zksync_fork(url):
     boa.set_env(ZksyncEnv.from_url(url))
     boa.env.fork()
@@ -20,5 +25,6 @@ def set_zksync_browser_env(address=None):
 
 
 boa.set_zksync_env = set_zksync_env
+boa.set_zksync_test_env = set_zksync_test_env
 boa.set_zksync_fork = set_zksync_fork
 boa.set_zksync_browser_env = set_zksync_browser_env

boa_zksync/compile.py

@@ -6,9 +6,12 @@
 from boa_zksync.types import ZksyncCompilerData
 
 
-def compile_zksync(filename: str, compiler_args=None) -> ZksyncCompilerData:
+def compile_zksync(
+    contract_name: str, filename: str, compiler_args=None, source_code=None
+) -> ZksyncCompilerData:
     vyper_path = which("vyper")
     assert vyper_path, "Vyper executable not found"
+    compiler_args = compiler_args or []
     compile_result = subprocess.run(
         [
             "zkvyper",
@@ -19,7 +22,7 @@ def compile_zksync(filename: str, compiler_args=None) -> ZksyncCompilerData:
             "-f",
             "combined_json",
             # pass any extra compiler args
-            *(compiler_args or []),
+            *compiler_args,
             # pass the file name
             "--",
             filename,
@@ -29,7 +32,12 @@ def compile_zksync(filename: str, compiler_args=None) -> ZksyncCompilerData:
 
     assert compile_result.returncode == 0, compile_result.stderr.decode()
     output = json.loads(compile_result.stdout.decode())
-    return ZksyncCompilerData(**output[filename])
+    if source_code is None:
+        with open(filename) as file:
+            source_code = file.read()
+    return ZksyncCompilerData(
+        contract_name, source_code, compiler_args, **output[filename]
+    )
 
 
 def compile_zksync_source(
@@ -39,4 +47,4 @@ def compile_zksync_source(
         filename = f"{tempdir}/{name}.vy"
         with open(filename, "w") as file:
             file.write(source_code)
-        return compile_zksync(filename, compiler_args)
+        return compile_zksync(name, filename, compiler_args, source_code)

boa_zksync/contract.py

@@ -0,0 +1,145 @@
+import textwrap
+from contextlib import contextmanager
+
+from boa.contracts.abi.abi_contract import ABIContract, ABIFunction
+from boa.contracts.vyper.compiler_utils import (
+    generate_source_for_internal_fn,
+    generate_source_for_arbitrary_stmt,
+    detect_statement_type,
+)
+from boa.rpc import to_bytes
+from cached_property import cached_property
+from vyper.semantics.analysis.base import VarInfo
+from vyper.semantics.types.function import ContractFunctionT
+
+from boa_zksync.compile import compile_zksync_source
+
+
+class ZksyncContract(ABIContract):
+    """
+    A contract deployed to the Zksync network.
+    """
+
+    def eval(self, code):
+        return ZksyncEval(code, self)()
+
+    @contextmanager
+    def override_vyper_namespace(self):
+        yield
+
+    @cached_property
+    def _storage(self):
+        def storage(): return None
+        for name, var in self.compiler_data.global_ctx.variables.items():
+            if not var.is_immutable and not var.is_constant:
+                setattr(storage, name, ZksyncInternalVariable(var, name, self))
+        return storage
+
+    @cached_property
+    def internal(self):
+        def internal(): return None
+        for fn in self.compiler_data.global_ctx.functions:
+            typ = fn._metadata["type"]
+            if typ.is_internal:
+                setattr(internal, fn.name, ZksyncInternalFunction(typ, self))
+        return internal
+
+
+class _ZksyncInternal(ABIFunction):
+    """
+    An ABI function that temporarily changes the bytecode at the contract's address.
+    """
+
+    @cached_property
+    def _override_bytecode(self):
+        data = self.contract.compiler_data
+        source = "\n".join((data.source_code, self.source_code))
+        compiled = compile_zksync_source(source, self.name, data.compiler_args)
+        return to_bytes(compiled.bytecode)
+
+    @property
+    def source_code(self):
+        raise NotImplementedError  # to be implemented in subclasses
+
+    def __call__(self, *args, **kwargs):
+        env = self.contract.env
+        env.set_code(self.contract.address, self._override_bytecode)
+        try:
+            return super().__call__(*args, **kwargs)
+        finally:
+            env.set_code(self.contract.address, to_bytes(self.contract._bytecode))
+
+
+class ZksyncInternalFunction(_ZksyncInternal):
+    def __init__(self, fn: ContractFunctionT, contract: ZksyncContract):
+        abi = {
+            "anonymous": False,
+            "inputs": [
+                {"name": arg.name, "type": arg.typ.abi_type.selector_name()}
+                for arg in fn.arguments
+            ],
+            "outputs": (
+                [{"name": fn.name, "type": fn.return_type.abi_type.selector_name()}]
+                if fn.return_type
+                else []
+            ),
+            "name": f"__boa_private_{fn.name}__",
+            "type": "function",
+        }
+        super().__init__(abi, contract._name)
+        self.contract = contract
+        self.func_t = fn
+
+    @cached_property
+    def source_code(self):
+        return generate_source_for_internal_fn(self)
+
+
+class ZksyncInternalVariable(_ZksyncInternal):
+    def __init__(self, var: VarInfo, name: str, contract: ZksyncContract):
+        abi = {
+            "anonymous": False,
+            "inputs": [],
+            "outputs": [{"name": name, "type": var.typ.abi_type.selector_name()}],
+            "name": f"__boa_private_{name}__",
+            "type": "function",
+        }
+        super().__init__(abi, contract._name)
+        self.contract = contract
+        self.var = var
+        self.var_name = name
+
+    def get(self):
+        return self.__call__()
+
+    @cached_property
+    def source_code(self):
+        return textwrap.dedent(
+            f"""
+            @external
+            @payable
+            def __boa_private_{self.var_name}__() -> {self.var.typ.abi_type.selector_name()}:
+                return self.{self.var_name}
+        """
+        )
+
+
+class ZksyncEval(_ZksyncInternal):
+    def __init__(self, code: str, contract: ZksyncContract):
+        typ = detect_statement_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)

boa_zksync/deployer.py

@@ -1,13 +1,15 @@
 from functools import cached_property
 
 from boa import Env
-from boa.contracts.abi.abi_contract import ABIContract, ABIContractFactory, ABIFunction
+from boa.contracts.abi.abi_contract import ABIContractFactory, ABIFunction
 from boa.rpc import to_bytes
 from boa.util.abi import Address
 
+from boa_zksync.contract import ZksyncContract
+
 
 class ZksyncDeployer(ABIContractFactory):
-    def deploy(self, *args, value=0, **kwargs):
+    def deploy(self, *args, value=0, **kwargs) -> ZksyncContract:
         env = Env.get_singleton()
         from boa_zksync.environment import ZksyncEnv
 
@@ -24,31 +26,38 @@ def deploy(self, *args, value=0, **kwargs):
                 else b""
             ),
         )
-        address = Address(address)
-        abi_contract = ABIContract(
-            self._name,
-            self.abi,
-            self._functions,
-            address=address,
-            filename=self._filename,
-            env=env,
-            compiler_data=self.compiler_data,
-        )
-        env.register_contract(address, abi_contract)
-        return abi_contract
+        return self.at(address)
 
-    def deploy_as_blueprint(self, *args, **kwargs):
+    def deploy_as_blueprint(self, *args, **kwargs) -> ZksyncContract:
         """
         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):
+    def constructor(self) -> ABIFunction:
         """
         Get the constructor function of the contract.
         :raises: StopIteration if the constructor is not found.
         """
         ctor_abi = next(i for i in self.abi if i["type"] == "constructor")
         return ABIFunction(ctor_abi, contract_name=self._name)
+
+    def at(self, address: Address | str) -> ZksyncContract:
+        """
+        Create an ABI contract object for a deployed contract at `address`.
+        """
+        address = Address(address)
+        env = Env.get_singleton()
+        contract = ZksyncContract(
+            self._name,
+            self.abi,
+            self._functions,
+            address=address,
+            filename=self._filename,
+            env=env,
+            compiler_data=self.compiler_data,
+        )
+        env.register_contract(address, contract)
+        return contract