refactor: simplify trait impl with ambassador

Cargo.toml

@@ -32,6 +32,7 @@ tokio = "1"
 
 [dependencies]
 ahash = { workspace = true }
+ambassador = "0.5"
 anes = "0.2"
 anyhow = { workspace = true }
 argon2 = "0.5"

src/beacon/drand.rs

@@ -1,6 +1,8 @@
 // Copyright 2019-2026 ChainSafe Systems
 // SPDX-License-Identifier: Apache-2.0, MIT
 
+#![allow(dead_code)]
+
 use std::sync::LazyLock;
 use std::time::Duration;
 use std::{borrow::Cow, num::NonZeroUsize};
@@ -16,13 +18,13 @@ use crate::shim::version::NetworkVersion;
 use crate::utils::cache::SizeTrackingLruCache;
 use crate::utils::misc::env::is_env_truthy;
 use crate::utils::net::global_http_client;
+use ambassador::{Delegate, delegatable_trait};
 use anyhow::Context as _;
 use backon::{ExponentialBuilder, Retryable};
 use bls_signatures::Serialize as _;
 use itertools::Itertools as _;
 use nonzero_ext::nonzero;
 use serde::{Deserialize as SerdeDeserialize, Serialize as SerdeSerialize};
-use spire_enum::prelude::delegated_enum;
 use tracing::debug;
 use url::Url;
 
@@ -140,38 +142,12 @@ impl BeaconSchedule {
     }
 }
 
-#[delegated_enum(impl_conversions)]
+#[derive(Delegate, derive_more::From)]
+#[delegate(Beacon)]
 pub enum BeaconImpl {
     Drand(DrandBeacon),
     #[cfg(test)]
-    Mock(super::mock_beacon::MockBeacon),
-}
-
-impl Beacon for BeaconImpl {
-    /// Gets the `drand` network
-    fn network(&self) -> DrandNetwork {
-        delegate_beacon_impl!(self.network())
-    }
-
-    /// Verify beacon entries that are sorted by round.
-    fn verify_entries(&self, entries: &[BeaconEntry], prev: &BeaconEntry) -> anyhow::Result<bool> {
-        delegate_beacon_impl!(self.verify_entries(entries, prev))
-    }
-
-    /// Returns a `BeaconEntry` given a round. It fetches the `BeaconEntry` from a `Drand` node over [`gRPC`](https://grpc.io/)
-    /// In the future, we will cache values, and support streaming.
-    async fn entry(&self, round: u64) -> anyhow::Result<BeaconEntry> {
-        delegate_beacon_impl!(self.entry(round).await)
-    }
-
-    /// Returns the most recent beacon round for the given Filecoin chain epoch.
-    fn max_beacon_round_for_epoch(
-        &self,
-        network_version: NetworkVersion,
-        fil_epoch: ChainEpoch,
-    ) -> u64 {
-        delegate_beacon_impl!(self.max_beacon_round_for_epoch(network_version, fil_epoch))
-    }
+    Mock(crate::beacon::mock_beacon::MockBeacon),
 }
 
 /// Contains height at which the beacon is activated, as well as the beacon
@@ -190,6 +166,7 @@ impl BeaconPoint {
 
 /// Trait used as the interface to be able to retrieve bytes from a randomness
 /// beacon.
+#[delegatable_trait]
 pub trait Beacon {
     /// Gets the `drand` network
     fn network(&self) -> DrandNetwork;

src/chain/store/base_fee.rs

@@ -4,7 +4,7 @@
 use std::sync::LazyLock;
 
 use crate::blocks::Tipset;
-use crate::message::Message;
+use crate::message::MessageReadWrite;
 use crate::shim::clock::ChainEpoch;
 use crate::shim::econ::{BLOCK_GAS_LIMIT, TokenAmount};
 use crate::utils::misc::env::env_or_default;
@@ -93,8 +93,8 @@ where
         let (bls_msgs, secp_msgs) = crate::chain::block_messages(db, b)?;
         for m in bls_msgs
             .iter()
-            .map(|m| m as &dyn Message)
-            .chain(secp_msgs.iter().map(|m| m as &dyn Message))
+            .map(|m| m as &dyn MessageReadWrite)
+            .chain(secp_msgs.iter().map(|m| m as &dyn MessageReadWrite))
         {
             if seen.insert((m.from(), m.sequence())) {
                 limits.push(m.gas_limit());

src/message/chain_message.rs

@@ -1,10 +1,11 @@
 // Copyright 2019-2026 ChainSafe Systems
 // SPDX-License-Identifier: Apache-2.0, MIT
 
-use super::{Message as MessageTrait, MessageRead};
+use super::*;
 use crate::message::signed_message::SignedMessage;
 use crate::shim::message::MethodNum;
 use crate::shim::{address::Address, econ::TokenAmount, message::Message};
+use ambassador::Delegate;
 use fvm_ipld_encoding::RawBytes;
 use get_size2::GetSize;
 use serde::{Deserialize, Serialize};
@@ -14,7 +15,10 @@ use std::sync::Arc;
 /// `Enum` to encapsulate signed and unsigned messages. Useful when working with
 /// both types
 #[delegated_enum]
-#[derive(Clone, Debug, Serialize, Deserialize, Hash, Eq, PartialEq, GetSize, derive_more::From)]
+#[derive(
+    Clone, Debug, Serialize, Deserialize, Hash, Eq, PartialEq, GetSize, derive_more::From, Delegate,
+)]
+#[delegate(MessageRead)]
 #[serde(untagged)]
 pub enum ChainMessage {
     Unsigned(Arc<Message>),
@@ -53,49 +57,9 @@ impl ChainMessage {
     pub fn equal_call(&self, other: &Self) -> bool {
         self.message().equal_call(other.message())
     }
-
-    pub fn set_sequence(&mut self, new_sequence: u64) {
-        match self {
-            Self::Unsigned(m) => Arc::make_mut(m).set_sequence(new_sequence),
-            Self::Signed(sm) => Arc::make_mut(sm).set_sequence(new_sequence),
-        }
-    }
-}
-
-impl MessageRead for ChainMessage {
-    fn from(&self) -> Address {
-        delegate_chain_message!(self.from())
-    }
-    fn to(&self) -> Address {
-        delegate_chain_message!(self.to())
-    }
-    fn sequence(&self) -> u64 {
-        delegate_chain_message!(self.sequence())
-    }
-    fn value(&self) -> TokenAmount {
-        delegate_chain_message!(self.value())
-    }
-    fn method_num(&self) -> MethodNum {
-        delegate_chain_message!(self.method_num())
-    }
-    fn params(&self) -> &RawBytes {
-        delegate_chain_message!(self.params())
-    }
-    fn gas_limit(&self) -> u64 {
-        delegate_chain_message!(self.gas_limit())
-    }
-    fn required_funds(&self) -> TokenAmount {
-        delegate_chain_message!(self.required_funds())
-    }
-    fn gas_fee_cap(&self) -> TokenAmount {
-        delegate_chain_message!(self.gas_fee_cap())
-    }
-    fn gas_premium(&self) -> TokenAmount {
-        delegate_chain_message!(self.gas_premium())
-    }
 }
 
-impl MessageTrait for ChainMessage {
+impl MessageReadWrite for ChainMessage {
     fn set_gas_limit(&mut self, amount: u64) {
         delegate_chain_message!(self => |i| Arc::make_mut(i).set_gas_limit(amount))
     }

src/message/mod.rs

@@ -5,8 +5,9 @@ pub mod chain_message;
 pub mod signed_message;
 
 use crate::shim::message::MethodNum;
-use crate::shim::{address::Address, econ::TokenAmount, message::Message as ShimMessage};
+use crate::shim::{address::Address, econ::TokenAmount, message::Message};
 use crate::shim::{gas::Gas, version::NetworkVersion};
+use ambassador::delegatable_trait;
 pub use chain_message::ChainMessage;
 use fvm_ipld_encoding::RawBytes;
 use num::Zero;
@@ -15,6 +16,7 @@ pub use signed_message::SignedMessage;
 /// Message interface to make read-only interactions with Signed and unsigned messages in a generic
 /// context.
 #[auto_impl::auto_impl(&, Arc)]
+#[delegatable_trait]
 pub trait MessageRead {
     /// Returns the from address of the message.
     fn from(&self) -> Address;
@@ -51,7 +53,7 @@ pub trait MessageRead {
 
 /// Message interface to interact with Signed and unsigned messages in a generic
 /// context.
-pub trait Message: MessageRead {
+pub trait MessageReadWrite: MessageRead {
     /// sets the gas limit for the message.
     fn set_gas_limit(&mut self, amount: u64);
     /// sets a new sequence to the message.
@@ -62,7 +64,7 @@ pub trait Message: MessageRead {
     fn set_gas_premium(&mut self, prem: TokenAmount);
 }
 
-impl MessageRead for ShimMessage {
+impl MessageRead for Message {
     fn from(&self) -> Address {
         self.from
     }
@@ -95,7 +97,7 @@ impl MessageRead for ShimMessage {
     }
 }
 
-impl Message for ShimMessage {
+impl MessageReadWrite for Message {
     fn set_gas_limit(&mut self, token_amount: u64) {
         self.gas_limit = token_amount;
     }
@@ -112,7 +114,7 @@ impl Message for ShimMessage {
 
 /// Semantic validation and validates the message has enough gas.
 pub fn valid_for_block_inclusion(
-    msg: &ShimMessage,
+    msg: &Message,
     min_gas: Gas,
     version: NetworkVersion,
 ) -> anyhow::Result<()> {
@@ -189,7 +191,7 @@ mod tests {
         .collect_vec();
 
         for (base_fee, gas_fee_cap, gas_premium, expected) in test_cases.into_iter() {
-            let msg = ShimMessage {
+            let msg = Message {
                 gas_fee_cap: gas_fee_cap.clone(),
                 gas_premium: gas_premium.clone(),
                 ..Default::default()

src/message/signed_message.rs

@@ -1,7 +1,7 @@
 // Copyright 2019-2026 ChainSafe Systems
 // SPDX-License-Identifier: Apache-2.0, MIT
 
-use super::{Message as MessageTrait, MessageRead};
+use super::{MessageRead, MessageReadWrite};
 use crate::eth::EthChainId;
 use crate::shim::message::MethodNum;
 use crate::shim::{
@@ -141,7 +141,7 @@ impl MessageRead for SignedMessage {
     }
 }
 
-impl MessageTrait for SignedMessage {
+impl MessageReadWrite for SignedMessage {
     fn set_gas_limit(&mut self, token_amount: u64) {
         self.message.set_gas_limit(token_amount);
     }

src/rpc/methods/eth.rs

@@ -26,7 +26,7 @@ use crate::eth::{
     EthLegacyHomesteadTxArgs, parse_eth_transaction,
 };
 use crate::lotus_json::{HasLotusJson, lotus_json_with_self};
-use crate::message::{ChainMessage, Message as _, MessageRead as _, SignedMessage};
+use crate::message::{ChainMessage, MessageRead as _, MessageReadWrite as _, SignedMessage};
 use crate::rpc::{
     ApiPaths, Ctx, EthEventHandler, LOOKBACK_NO_LIMIT, Permission, RpcMethod, RpcMethodExt as _,
     error::ServerError,

src/rpc/methods/gas.rs

@@ -4,7 +4,7 @@
 use super::state::InvocResult;
 use crate::blocks::Tipset;
 use crate::chain::{BASE_FEE_MAX_CHANGE_DENOM, BLOCK_GAS_TARGET};
-use crate::message::{ChainMessage, Message as _, MessageRead as _, SignedMessage};
+use crate::message::{ChainMessage, MessageRead as _, MessageReadWrite as _, SignedMessage};
 use crate::rpc::chain::FlattenedApiMessage;
 use crate::rpc::{ApiPaths, Ctx, Permission, RpcMethod, error::ServerError, types::*};
 use crate::shim::executor::ApplyRet;

src/state_manager/mod.rs

@@ -24,7 +24,7 @@ use crate::interpreter::{
 };
 use crate::interpreter::{MessageCallbackCtx, VMTrace};
 use crate::lotus_json::{LotusJson, lotus_json_with_self};
-use crate::message::{ChainMessage, Message as _, MessageRead as _, SignedMessage};
+use crate::message::{ChainMessage, MessageRead as _, MessageReadWrite as _, SignedMessage};
 use crate::networks::ChainConfig;
 use crate::rpc::state::{ApiInvocResult, InvocResult, MessageGasCost};
 use crate::rpc::types::{MiningBaseInfo, SectorOnChainInfo};