ABI dir fixes

beamer/agent/agent.py

@@ -3,7 +3,7 @@
 from itertools import permutations
 
 import structlog
-from eth_typing import Address, BlockNumber, ChecksumAddress
+from eth_typing import Address, ChecksumAddress
 from web3.middleware import latest_block_based_cache_middleware
 
 import beamer.agent.metrics
@@ -12,31 +12,19 @@
 from beamer.agent.state_machine import Context
 from beamer.agent.tracker import Tracker
 from beamer.agent.util import BaseChain, Chain
-from beamer.contracts import ContractInfo, make_contracts
+from beamer.contracts import ABIManager, obtain_contract
 from beamer.typing import ChainId, TransferDirection
 from beamer.util import make_web3
 
 log = structlog.get_logger(__name__)
 
 
-def _get_contracts_info(config: Config, chain_id: ChainId) -> dict[str, ContractInfo]:
-    info = config.deployment_info.get(chain_id)
-    if info is None:
-        raise RuntimeError(f"Deployment info for chain ID {chain_id} not available")
-    return info
-
-
-def _get_deployment_block(contract_info: dict[str, ContractInfo]) -> BlockNumber:
-    request_manager_deployment_block = contract_info["RequestManager"].deployment_block
-    fill_manager_deployment_block = contract_info["FillManager"].deployment_block
-    return min(request_manager_deployment_block, fill_manager_deployment_block)
-
-
 class Agent:
     def __init__(self, config: Config):
         self._config = config
         self._stopped = threading.Event()
         self._stopped.set()
+        self._abi_manager = ABIManager(config.abi_dir)
         self._init()
 
     def _init_l1_chain(self) -> BaseChain:
@@ -51,14 +39,19 @@ def _init_chains(self) -> dict[ChainId, Chain]:
             chain_id = ChainId(w3.eth.chain_id)
             if chain_id in chains:
                 continue
-            contracts_info = _get_contracts_info(self._config, chain_id)
-            contracts = make_contracts(w3, contracts_info)
-            request_manager = contracts["RequestManager"]
-            fill_manager = contracts["FillManager"]
+
+            deployment = beamer.artifacts.load(self._config.artifacts_dir, chain_id)
+            if deployment is None:
+                raise RuntimeError(f"Deployment artifact for chain ID {chain_id} not available")
+
+            assert deployment.chain is not None
+            request_manager = obtain_contract(w3, self._abi_manager, deployment, "RequestManager")
+            fill_manager = obtain_contract(w3, self._abi_manager, deployment, "FillManager")
+
             self._event_monitors[chain_id] = EventMonitor(
                 web3=w3,
                 contracts=(request_manager, fill_manager),
-                deployment_block=_get_deployment_block(contracts_info),
+                deployment_block=deployment.earliest_block,
                 poll_period=chain_config.poll_period,
                 confirmation_blocks=chain_config.confirmation_blocks,
                 on_new_events=[],
@@ -70,7 +63,8 @@ def _init_chains(self) -> dict[ChainId, Chain]:
                 id=chain_id,
                 name=chain_name,
                 tokens=self._config.token_checker.get_tokens_for_chain(chain_id),
-                contracts=contracts,
+                request_manager=request_manager,
+                fill_manager=fill_manager,
             )
         return chains
 

beamer/agent/config.py

@@ -9,7 +9,6 @@
 from eth_utils import to_wei
 
 from beamer.agent.util import TokenChecker
-from beamer.contracts import DeploymentInfo, load_deployment_info
 from beamer.typing import URL
 from beamer.util import account_from_keyfile
 
@@ -29,7 +28,8 @@ class ChainConfig:
 @dataclass
 class Config:
     account: LocalAccount
-    deployment_info: DeploymentInfo
+    abi_dir: Path
+    artifacts_dir: Path
     base_chain_rpc_url: URL
     token_checker: TokenChecker
     fill_wait_time: int
@@ -139,12 +139,12 @@ def load(config_path: Path, options: dict[str, Any]) -> Config:
     password = _get_value(config, "account.password")
     account = account_from_keyfile(path, password)
 
-    deployment_info = load_deployment_info(Path(config["artifacts-dir"]), Path(config["abi-dir"]))
     token_checker = TokenChecker(list(config["tokens"].values()))
 
     return Config(
         account=account,
-        deployment_info=deployment_info,
+        abi_dir=Path(config["abi-dir"]),
+        artifacts_dir=Path(config["artifacts-dir"]),
         token_checker=token_checker,
         base_chain_rpc_url=config["base-chain"]["rpc-url"],
         fill_wait_time=config["fill-wait-time"],

beamer/agent/util.py

@@ -26,15 +26,8 @@ def rpc_url(self) -> URL:
 class Chain(BaseChain):
     name: str
     tokens: list[tuple[ChainId, ChecksumAddress]]
-    contracts: dict[str, Contract]
-
-    @property
-    def request_manager(self) -> Contract:
-        return self.contracts["RequestManager"]
-
-    @property
-    def fill_manager(self) -> Contract:
-        return self.contracts["FillManager"]
+    request_manager: Contract
+    fill_manager: Contract
 
 
 @dataclass(frozen=True)

beamer/artifacts.py

@@ -0,0 +1,72 @@
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import apischema
+from apischema import properties, schema
+from apischema.metadata import none_as_undefined, validators
+from eth_typing import ChecksumAddress
+from eth_utils import is_checksum_address
+
+from beamer.typing import BlockNumber, ChainId
+
+
+class ValidationError(Exception):
+    def __str__(self) -> str:
+        assert isinstance(self.__cause__, apischema.ValidationError)
+        return "\n".join(map(str, self.__cause__.errors))  # pylint: disable=no-member
+
+
+def _validate_address(address: str) -> None:
+    if not is_checksum_address(address):
+        raise apischema.ValidationError(f"expected a checksum address: {address}")
+
+
+@dataclass
+class DeployedContractInfo:
+    beamer_commit: str
+    tx_hash: str
+    address: ChecksumAddress = field(metadata=validators(_validate_address))
+    deployment_block: BlockNumber = field(metadata=schema(min=0))
+    deployment_args: list[str | int]
+
+
+@dataclass(frozen=True)
+class ChainDeployment:
+    chain_id: ChainId = field(metadata=schema(min=1))
+    contracts: dict[str, DeployedContractInfo] = field(metadata=properties)
+
+
+@dataclass(frozen=True)
+class Deployment:
+    deployer: ChecksumAddress
+    base: ChainDeployment
+    chain: ChainDeployment | None = field(default=None, metadata=none_as_undefined)
+
+    @staticmethod
+    def from_file(artifact: Path) -> "Deployment":
+        with open(artifact, "rt") as f:
+            data = json.load(f)
+
+        try:
+            deployment = apischema.deserialize(Deployment, data)
+        except apischema.ValidationError as exc:
+            raise ValidationError from exc
+        return deployment
+
+    def to_file(self, artifact: Path) -> None:
+        with open(artifact, "wt") as f:
+            json.dump(apischema.serialize(self), f, indent=4)
+
+    @property
+    def earliest_block(self) -> BlockNumber:
+        assert self.chain is not None
+        return min(info.deployment_block for info in self.chain.contracts.values())
+
+
+def load(artifacts_dir: Path, chain_id: ChainId) -> Deployment:
+    path = next(artifacts_dir.glob(f"{chain_id}-*.deployment.json"))
+    deployment = Deployment.from_file(path)
+    chain = deployment.chain or deployment.base
+    assert chain.chain_id == chain_id
+    return deployment

beamer/config/commands.py

@@ -8,10 +8,10 @@
 from web3 import Web3
 
 import beamer.contracts
-import beamer.deploy.config
 import beamer.util
+from beamer.artifacts import Deployment
 from beamer.config.state import ChainConfig, Configuration, DesiredConfiguration, TokenConfig
-from beamer.deploy.artifacts import Deployment
+from beamer.contracts import ABIManager, obtain_contract
 from beamer.events import (
     ChainUpdated,
     Event,
@@ -122,7 +122,7 @@ def read(
     """Read latest contract configuration state from the chain and store it into STATE_PATH."""
     beamer.util.setup_logging(log_level="DEBUG", log_json=False)
 
-    rpc_info = beamer.deploy.config.load_rpc_info(rpc_file)
+    rpc_info = beamer.util.load_rpc_info(rpc_file)
     deployment = Deployment.from_file(artifact)
 
     assert deployment.chain is not None
@@ -132,8 +132,9 @@ def read(
     assert w3.eth.chain_id == chain_id
     log.info("Connected to RPC", url=url)
 
-    request_manager = deployment.obtain_contract(w3, "chain", "RequestManager")
-    fill_manager = deployment.obtain_contract(w3, "chain", "FillManager")
+    abi_manager = ABIManager(abi_dir)
+    request_manager = obtain_contract(w3, abi_manager, deployment, "RequestManager")
+    fill_manager = obtain_contract(w3, abi_manager, deployment, "FillManager")
 
     if state_path.exists():
         config = Configuration.from_file(state_path)
@@ -283,10 +284,10 @@ def _ensure_same_tokens_have_same_addresses(
 
 
 def _ensure_no_config_updates_since(
-    w3: Web3, deployment: Deployment, start_block: BlockNumber
+    w3: Web3, abi_manager: ABIManager, deployment: Deployment, start_block: BlockNumber
 ) -> None:
-    request_manager = deployment.obtain_contract(w3, "chain", "RequestManager")
-    fill_manager = deployment.obtain_contract(w3, "chain", "FillManager")
+    request_manager = obtain_contract(w3, abi_manager, deployment, "RequestManager")
+    fill_manager = obtain_contract(w3, abi_manager, deployment, "FillManager")
 
     fetcher = EventFetcher(
         w3, (request_manager, fill_manager), start_block=start_block, confirmation_blocks=0
@@ -378,7 +379,7 @@ def write(
     account = beamer.util.account_from_keyfile(keystore_file, password)
     log.info("Loaded keystore file", address=account.address)
 
-    rpc_info = beamer.deploy.config.load_rpc_info(rpc_file)
+    rpc_info = beamer.util.load_rpc_info(rpc_file)
     deployment = Deployment.from_file(artifact)
 
     assert deployment.chain is not None
@@ -388,6 +389,7 @@ def write(
     assert w3.eth.chain_id == chain_id
     log.info("Connected to RPC", url=url)
 
+    abi_manager = ABIManager(abi_dir)
     current_config = Configuration.from_file(current_state_path)
     desired_config = DesiredConfiguration.from_file(desired_state_path)
 
@@ -400,11 +402,11 @@ def write(
     # This means we need to make sure that there were no config updates
     # in block range [current_config.block + 1, latest_block].
     start_block = BlockNumber(current_config.block + 1)
-    _ensure_no_config_updates_since(w3, deployment, start_block)
+    _ensure_no_config_updates_since(w3, abi_manager, deployment, start_block)
 
     for contract, function, *args in _generate_updates(current_config, desired_config):
         log.info("Sending transaction", call=f"{contract}.{function}({', '.join(map(str, args))})")
-        contract = deployment.obtain_contract(w3, "chain", contract)
+        contract = obtain_contract(w3, abi_manager, deployment, contract)
         call = getattr(contract.functions, function)(*args)
         try:
             receipt = beamer.util.transact(call)

beamer/contracts.py

@@ -1,69 +1,57 @@
 import json
-from dataclasses import dataclass
+from collections import namedtuple
 from pathlib import Path
 from typing import cast
 
-import web3
 from web3 import Web3
 from web3.contract import Contract
 
-import beamer.deploy.artifacts
-from beamer.typing import BlockNumber, ChainId, ChecksumAddress
+import beamer.artifacts
 
 
-@dataclass
-class ContractInfo:
-    address: ChecksumAddress
-    deployment_block: BlockNumber
-    abi: list
+class ABIManager:
+    _CacheEntry = namedtuple("_CacheEntry", ("abi", "bytecode"))
 
+    def __init__(self, abi_dir: Path):
+        self.abi_dir = abi_dir
+        self._cache: dict[str, ABIManager._CacheEntry] = {}
 
-def make_contracts(w3: web3.Web3, contracts_info: dict[str, ContractInfo]) -> dict[str, Contract]:
-    return {
-        name: cast(Contract, w3.eth.contract(info.address, abi=info.abi, decode_tuples=True))
-        for name, info in contracts_info.items()
-    }
+    def get_abi(self, name: str) -> str:
+        entry = self._cache.get(name)
+        if entry is None:
+            entry = self._load_entry(name)
+            self._cache[name] = entry
+        return entry.abi
 
+    def get_bytecode(self, name: str) -> str:
+        entry = self._cache.get(name)
+        if entry is None:
+            entry = self._load_entry(name)
+            self._cache[name] = entry
+        return entry.bytecode
 
-def load_contract_abi(abi_dir: Path, contract_name: str) -> list:
-    with abi_dir.joinpath(f"{contract_name}.json").open("rt") as f:
-        data = json.load(f)
-    return data["abi"]
-
-
-DeploymentInfo = dict[ChainId, dict[str, ContractInfo]]
-
-
-def prepare_deployment_infos(
-    abi_dir: Path, contracts: dict[str, beamer.deploy.artifacts.DeployedContractInfo]
-) -> dict[str, ContractInfo]:
-    abis = {}
-    infos = {}
-    for name, contract in contracts.items():
-        if name not in abis:
-            abis[name] = load_contract_abi(abi_dir, name)
-        abi = abis[name]
-        infos[name] = ContractInfo(
-            address=contract.address,
-            deployment_block=contract.deployment_block,
-            abi=abi,
+    def _load_entry(self, name: str) -> _CacheEntry:
+        path = self.abi_dir.joinpath(f"{name}.json")
+        with path.open("rt") as f:
+            data = json.load(f)
+        return ABIManager._CacheEntry(
+            abi=data["abi"], bytecode=data["runtimeBytecode"]["bytecode"]
         )
-    return infos
 
 
-def load_deployment_info(artifacts_dir: Path, abi_dir: Path) -> DeploymentInfo:
-    deployment_info = {}
-    for artifact_path in artifacts_dir.glob("*.deployment.json"):
-        deployment = beamer.deploy.artifacts.Deployment.from_file(artifact_path)
-        if deployment.chain is None:
-            continue
-        deployment_info[deployment.chain.chain_id] = prepare_deployment_infos(
-            abi_dir, deployment.chain.contracts
-        )
-    return deployment_info
+def obtain_contract(
+    w3: Web3, abi_manager: ABIManager, deployment: beamer.artifacts.Deployment, name: str
+) -> Contract:
+    chain_id = w3.eth.chain_id
 
+    if chain_id == deployment.base.chain_id and name in deployment.base.contracts:
+        address = deployment.base.contracts[name].address
+    elif deployment.chain is not None:
+        if chain_id == deployment.chain.chain_id and name in deployment.chain.contracts:
+            address = deployment.chain.contracts[name].address
+    else:
+        raise ValueError(f"{name} not found on chain with ID {chain_id} in {deployment}")
 
-def contracts_for_web3(web3: Web3, artifacts_dir: Path, abi_dir: Path) -> dict[str, Contract]:
-    deployment_info = load_deployment_info(artifacts_dir, abi_dir)
-    chain_id = ChainId(web3.eth.chain_id)
-    return make_contracts(web3, deployment_info[chain_id])
+    abi = abi_manager.get_abi(name)
+    contract = w3.eth.contract(address, abi=abi, decode_tuples=True)
+    return cast(Contract, contract)