feat: add a trait for call-like container enums

Author: DaniPopes

crates/sol-macro/src/expand/contract.rs

@@ -40,7 +40,7 @@ pub(super) fn expand(cx: &ExpCtxt<'_>, contract: &ItemContract) -> Result<TokenS
 
     let functions_enum = if functions.len() > 1 {
         let mut attrs = d_attrs.clone();
-        let doc_str = format!("Container for all the [`{name}`] function calls.");
+        let doc_str = format!("Container for all the `{name}` function calls.");
         attrs.push(parse_quote!(#[doc = #doc_str]));
         Some(expand_functions_enum(cx, name, functions, &attrs))
     } else {
@@ -49,7 +49,7 @@ pub(super) fn expand(cx: &ExpCtxt<'_>, contract: &ItemContract) -> Result<TokenS
 
     let errors_enum = if errors.len() > 1 {
         let mut attrs = d_attrs;
-        let doc_str = format!("Container for all the [`{name}`] custom errors.");
+        let doc_str = format!("Container for all the `{name}` custom errors.");
         attrs.push(parse_quote!(#[doc = #doc_str]));
         Some(expand_errors_enum(cx, name, errors, &attrs))
     } else {
@@ -71,11 +71,11 @@ pub(super) fn expand(cx: &ExpCtxt<'_>, contract: &ItemContract) -> Result<TokenS
 
 fn expand_functions_enum(
     cx: &ExpCtxt<'_>,
-    name: &SolIdent,
+    contract_name: &SolIdent,
     functions: Vec<&ItemFunction>,
     attrs: &[Attribute],
 ) -> TokenStream {
-    let name = format_ident!("{name}Calls");
+    let name = format_ident!("{contract_name}Calls");
     let variants: Vec<_> = functions
         .iter()
         .map(|f| cx.function_name_ident(f).0)
@@ -86,24 +86,24 @@ fn expand_functions_enum(
     let min_data_len = functions
         .iter()
         .map(|function| r#type::params_min_data_size(cx, &function.arguments))
-        .max()
+        .min()
         .unwrap();
     let trt = Ident::new("SolCall", Span::call_site());
     expand_call_like_enum(name, &variants, &types, min_data_len, trt, attrs)
 }
 
 fn expand_errors_enum(
     cx: &ExpCtxt<'_>,
-    name: &SolIdent,
+    contract_name: &SolIdent,
     errors: Vec<&ItemError>,
     attrs: &[Attribute],
 ) -> TokenStream {
-    let name = format_ident!("{name}Errors");
+    let name = format_ident!("{contract_name}Errors");
     let variants: Vec<_> = errors.iter().map(|error| error.name.0.clone()).collect();
     let min_data_len = errors
         .iter()
         .map(|error| r#type::params_min_data_size(cx, &error.parameters))
-        .max()
+        .min()
         .unwrap();
     let trt = Ident::new("SolError", Span::call_site());
     expand_call_like_enum(name, &variants, &variants, min_data_len, trt, attrs)
@@ -120,64 +120,75 @@ fn expand_call_like_enum(
     assert_eq!(variants.len(), types.len());
     let name_s = name.to_string();
     let count = variants.len();
-    let min_data_len = min_data_len.min(4);
     quote! {
         #(#attrs)*
         pub enum #name {
             #(#variants(#types),)*
         }
 
-        // TODO: Implement these functions using traits?
         #[automatically_derived]
-        impl #name {
-            /// The number of variants.
-            pub const COUNT: usize = #count;
-
-            // no decode_raw is possible because we need the selector to know which variant to
-            // decode into
-
-            /// ABI-decodes the given data into one of the variants of `self`.
-            pub fn decode(data: &[u8], validate: bool) -> ::alloy_sol_types::Result<Self> {
-                if data.len() >= #min_data_len {
-                    // TODO: Replace with `data.split_array_ref` once it's stable
-                    let (selector, data) = data.split_at(4);
-                    let selector: &[u8; 4] =
-                        ::core::convert::TryInto::try_into(selector).expect("unreachable");
-                    match *selector {
-                        #(<#types as ::alloy_sol_types::#trt>::SELECTOR => {
-                            return <#types as ::alloy_sol_types::#trt>::decode_raw(data, validate)
-                                .map(Self::#variants)
-                        })*
-                        _ => {}
-                    }
+        impl ::alloy_sol_types::SolCalls for #name {
+            const NAME: &'static str = #name_s;
+            const MIN_DATA_LENGTH: usize = #min_data_len;
+            const COUNT: usize = #count;
+
+            #[inline]
+            fn selector(&self) -> [u8; 4] {
+                match self {#(
+                    Self::#variants(_) => <#types as ::alloy_sol_types::#trt>::SELECTOR,
+                )*}
+            }
+
+            #[inline]
+            fn type_check(selector: [u8; 4]) -> ::alloy_sol_types::Result<()> {
+                match selector {
+                    #(<#types as ::alloy_sol_types::#trt>::SELECTOR)|* => Ok(()),
+                    s => ::core::result::Result::Err(::alloy_sol_types::Error::unknown_selector(
+                        Self::NAME,
+                        s,
+                    )),
                 }
-                ::core::result::Result::Err(::alloy_sol_types::Error::type_check_fail(
-                    data,
-                    #name_s,
-                ))
             }
 
-            /// ABI-encodes `self` into the given buffer.
-            pub fn encode_raw(&self, out: &mut Vec<u8>) {
+            #[inline]
+            fn decode_raw(
+                selector: [u8; 4],
+                data: &[u8],
+                validate: bool
+            )-> ::alloy_sol_types::Result<Self> {
+                match selector {
+                    #(<#types as ::alloy_sol_types::#trt>::SELECTOR => {
+                        <#types as ::alloy_sol_types::#trt>::decode_raw(data, validate)
+                            .map(Self::#variants)
+                    })*
+                    s => ::core::result::Result::Err(::alloy_sol_types::Error::unknown_selector(
+                        Self::NAME,
+                        s,
+                    )),
+                }
+            }
+
+            #[inline]
+            fn encoded_size(&self) -> usize {
                 match self {#(
                     Self::#variants(inner) =>
-                        <#types as ::alloy_sol_types::#trt>::encode_raw(inner, out),
+                        <#types as ::alloy_sol_types::#trt>::encoded_size(inner),
                 )*}
             }
 
-            /// ABI-encodes `self` into the given buffer.
             #[inline]
-            pub fn encode(&self) -> Vec<u8> {
+            fn encode_raw(&self, out: &mut Vec<u8>) {
                 match self {#(
                     Self::#variants(inner) =>
-                        <#types as ::alloy_sol_types::#trt>::encode(inner),
+                        <#types as ::alloy_sol_types::#trt>::encode_raw(inner, out),
                 )*}
             }
         }
 
         #(
             #[automatically_derived]
             impl From<#types> for #name {
+                #[inline]
                 fn from(value: #types) -> Self {
                     Self::#variants(value)
                 }

crates/sol-types/src/errors.rs

@@ -40,6 +40,14 @@ pub enum Error {
         max: u8,
     },
 
+    /// Unknown selector.
+    UnknownSelector {
+        /// The type name.
+        name: &'static str,
+        /// The unknown selector.
+        selector: alloy_primitives::FixedBytes<4>,
+    },
+
     /// Hex error.
     FromHexError(hex::FromHexError),
 
@@ -73,6 +81,9 @@ impl fmt::Display for Error {
                 f,
                 "`{value}` is not a valid {name} enum value (max: `{max}`)"
             ),
+            Self::UnknownSelector { name, selector } => {
+                write!(f, "Unknown selector `{selector}` for {name}")
+            }
             Self::FromHexError(e) => e.fmt(f),
             Self::Other(e) => f.write_str(e),
         }
@@ -104,6 +115,15 @@ impl Error {
             data: hex::encode(data),
         }
     }
+
+    /// Instantiates a [`Error::UnknownSelector`] with the provided data.
+    #[inline]
+    pub fn unknown_selector(name: &'static str, selector: [u8; 4]) -> Self {
+        Self::UnknownSelector {
+            name,
+            selector: selector.into(),
+        }
+    }
 }
 
 impl From<hex::FromHexError> for Error {

crates/sol-types/src/lib.rs

@@ -181,8 +181,8 @@ mod impl_core;
 
 mod types;
 pub use types::{
-    data_type as sol_data, Encodable, EventTopic, Panic, PanicKind, Revert, SolCall, SolError,
-    SolEvent, SolStruct, SolType, TopicList,
+    data_type as sol_data, ContractError, Encodable, EventTopic, Panic, PanicKind, Revert, SolCall,
+    SolCalls, SolEnum, SolError, SolEvent, SolStruct, SolType, TopicList,
 };
 
 mod util;

crates/sol-types/src/types/calls.rs

@@ -0,0 +1,130 @@
+use crate::{Panic, Result, Revert, SolError};
+use alloc::vec::Vec;
+
+/// A collection of ABI-encoded call-like types. This currently includes
+/// [`SolCall`] and [`SolError`].
+///
+/// [`SolCall`]: crate::SolCall
+/// [`SolError`]: crate::SolError
+///
+/// ### Implementer's Guide
+///
+/// We do not recommend implementing this trait directly. Instead, we recommend
+/// using the [`sol`][crate::sol] proc macro to parse a Solidity contract
+/// definition.
+pub trait SolCalls: Sized {
+    /// The name of this type.
+    const NAME: &'static str;
+
+    /// The minimum length of the data for this type.
+    ///
+    /// This does *not* include the selector's length (4).
+    const MIN_DATA_LENGTH: usize;
+
+    /// The number of variants.
+    const COUNT: usize;
+
+    /// The selector of this type.
+    fn selector(&self) -> [u8; 4];
+
+    /// Checks if the given selector is known to this type.
+    fn type_check(selector: [u8; 4]) -> Result<()>;
+
+    /// ABI-decodes the given data into one of the variants of `self`.
+    fn decode_raw(selector: [u8; 4], data: &[u8], validate: bool) -> Result<Self>;
+
+    /// The size of the encoded data, *without* any selectors.
+    fn encoded_size(&self) -> usize;
+
+    /// ABI-encodes `self` into the given buffer, *without* any selectors.
+    fn encode_raw(&self, out: &mut Vec<u8>);
+
+    /// ABI-encodes `self` into the given buffer.
+    fn encode(&self) -> Vec<u8> {
+        let mut out = Vec::with_capacity(4 + self.encoded_size());
+        out.extend(self.selector());
+        self.encode_raw(&mut out);
+        out
+    }
+
+    /// ABI-decodes the given data into one of the variants of `self`.
+    #[inline]
+    fn decode(data: &[u8], validate: bool) -> Result<Self> {
+        if data.len() < Self::MIN_DATA_LENGTH + 4 {
+            Err(crate::Error::type_check_fail(data, Self::NAME))
+        } else {
+            let (selector, data) = crate::impl_core::split_array_ref(data);
+            Self::decode_raw(*selector, data, validate)
+        }
+    }
+}
+
+/// A generic contract error.
+///
+/// Contains a [`Revert`] or [`Panic`] error, or a custom error.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum ContractError<T> {
+    /// A contract's custom error.
+    CustomError(T),
+    /// A generic revert. See [`Revert`] for more information.
+    Revert(Revert),
+    /// A panic. See [`Panic`] for more information.
+    Panic(Panic),
+}
+
+impl<T: SolCalls> SolCalls for ContractError<T> {
+    const NAME: &'static str = "ContractError";
+
+    // revert is 64, panic is 32
+    const MIN_DATA_LENGTH: usize = if T::MIN_DATA_LENGTH < 32 {
+        T::MIN_DATA_LENGTH
+    } else {
+        32
+    };
+
+    const COUNT: usize = T::COUNT + 2;
+
+    #[inline]
+    fn selector(&self) -> [u8; 4] {
+        match self {
+            Self::CustomError(error) => error.selector(),
+            Self::Panic(_) => Panic::SELECTOR,
+            Self::Revert(_) => Revert::SELECTOR,
+        }
+    }
+
+    #[inline]
+    fn type_check(selector: [u8; 4]) -> Result<()> {
+        match selector {
+            Revert::SELECTOR | Panic::SELECTOR => Ok(()),
+            s => T::type_check(s),
+        }
+    }
+
+    #[inline]
+    fn decode_raw(selector: [u8; 4], data: &[u8], validate: bool) -> Result<Self> {
+        match selector {
+            Revert::SELECTOR => Revert::decode_raw(data, validate).map(Self::Revert),
+            Panic::SELECTOR => Panic::decode_raw(data, validate).map(Self::Panic),
+            _ => T::decode(data, validate).map(Self::CustomError),
+        }
+    }
+
+    #[inline]
+    fn encoded_size(&self) -> usize {
+        match self {
+            Self::CustomError(error) => error.encoded_size(),
+            Self::Panic(panic) => panic.encoded_size(),
+            Self::Revert(revert) => revert.encoded_size(),
+        }
+    }
+
+    #[inline]
+    fn encode_raw(&self, out: &mut Vec<u8>) {
+        match self {
+            Self::CustomError(error) => error.encode_raw(out),
+            Self::Panic(panic) => panic.encode_raw(out),
+            Self::Revert(revert) => revert.encode_raw(out),
+        }
+    }
+}

crates/sol-types/src/types/enum.rs

@@ -0,0 +1,36 @@
+use crate::{token::WordToken, Result, SolType, Word};
+use alloc::vec::Vec;
+
+/// Solidity enum. This is always a wrapper around a [`u8`].
+///
+/// ### Implementer's Guide
+///
+/// We do not recommend implementing this trait directly. Instead, we recommend
+/// using the [`sol`][crate::sol] proc macro to parse a Solidity error
+/// definition.
+pub trait SolEnum: Sized + Copy + Into<u8> + TryFrom<u8, Error = crate::Error> {
+    /// Tokenize the enum.
+    #[inline]
+    fn tokenize(self) -> WordToken {
+        WordToken(Word::with_last_byte(self.into()))
+    }
+
+    /// ABI decode the enum from the given buffer.
+    #[inline]
+    fn decode(data: &[u8], validate: bool) -> Result<Self> {
+        <crate::sol_data::Uint<8> as SolType>::decode_single(data, validate)
+            .and_then(Self::try_from)
+    }
+
+    /// ABI encode the enum into the given buffer.
+    #[inline]
+    fn encode_raw(self, out: &mut Vec<u8>) {
+        out.extend(self.tokenize().0)
+    }
+
+    /// ABI encode the enum.
+    #[inline]
+    fn encode(self) -> Vec<u8> {
+        self.tokenize().0.to_vec()
+    }
+}

crates/sol-types/src/types/mod.rs

@@ -1,3 +1,6 @@
+mod calls;
+pub use calls::{ContractError, SolCalls};
+
 pub mod data_type;
 
 mod r#enum;