feat: Add function handler support

MrWad3r · Rexagon · commit 5db329cb4138 · 2025-04-03T13:02:49.000+02:00
diff --git a/abi-proc/src/common.rs b/abi-proc/src/common.rs
@@ -1,35 +1,96 @@
-use proc_macro2::Ident;
+use proc_macro2::{Ident, Span};
 use quote::format_ident;
-use syn::{Attribute, LitStr, Path};
+use syn::spanned::Spanned;
+use syn::{Attribute, Error, LitStr, Path};
+
+pub enum HandlerType {
+    AbiType,
+    IntoAbi,
+    FromAbi,
+}
 
 pub struct FieldAttributes {
     pub custom_name: Option<Ident>,
-    pub custom_handler: Option<Path>,
+    pub mod_handler: Option<Path>,
+    pub with_handlers: Vec<(HandlerType, Path)>,
+}
+
+impl FieldAttributes {
+    pub fn check(&self) -> Result<(), Error> {
+        if !self.with_handlers.is_empty() && self.mod_handler.is_some() {
+            return Err(Error::new(
+                Span::call_site(),
+                "`with` attribute should not be used with other handling attributes",
+            ));
+        }
+        Ok(())
+    }
 }
 
-pub fn extract_field_attributes(attrs: &[Attribute]) -> FieldAttributes {
+const ABI: &str = "abi";
+const NAME: &str = "name";
+const WITH: &str = "with";
+const ABI_TYPE_WITH: &str = "abi_type_with";
+const INTO_ABI_WITH: &str = "into_abi_with";
+const FROM_ABI_WITH: &str = "from_abi_with";
+
+pub fn extract_field_attributes(attrs: &[Attribute]) -> Result<FieldAttributes, Error> {
     let mut attributes = FieldAttributes {
         custom_name: None,
-        custom_handler: None,
+        mod_handler: None,
+        with_handlers: Vec::new(),
     };
+
+    fn parse_path(meta: syn::meta::ParseNestedMeta) -> Result<Path, Error> {
+        let value = meta.value()?;
+        let path_str: syn::LitStr = value.parse()?;
+        let path = syn::parse_str::<Path>(&path_str.value())?;
+        Ok(path)
+    }
+
     for attr in attrs {
-        if !attr.path().is_ident("abi") {
+        if !attr.path().is_ident(ABI) {
             continue;
         }
-        let _ = attr.parse_nested_meta(|meta| {
-            if meta.path.is_ident("name") {
-                let value = meta.value()?;
-                let lit: LitStr = value.parse()?;
-                attributes.custom_name = Some(format_ident!("{}", lit.value(), span = lit.span()));
-            } else if meta.path.is_ident("with") {
-                let value = meta.value()?;
-                let path_str: syn::LitStr = value.parse()?;
-                let path = syn::parse_str::<Path>(&path_str.value())?;
-                attributes.custom_handler = Some(path);
+        attr.parse_nested_meta(|meta| {
+            match &meta.path {
+                path if path.is_ident(NAME) => {
+                    let value = meta.value()?;
+                    let lit: LitStr = value.parse()?;
+                    attributes.custom_name =
+                        Some(format_ident!("{}", lit.value(), span = lit.span()));
+                }
+                path if path.is_ident(WITH) => {
+                    let value = meta.value()?;
+                    let path_str: syn::LitStr = value.parse()?;
+                    let path = syn::parse_str::<Path>(&path_str.value())?;
+                    attributes.mod_handler = Some(path);
+                }
+                path if path.is_ident(ABI_TYPE_WITH) => {
+                    let path = parse_path(meta)?;
+                    attributes.with_handlers.push((HandlerType::AbiType, path));
+                }
+                path if path.is_ident(INTO_ABI_WITH) => {
+                    let path = parse_path(meta)?;
+                    attributes.with_handlers.push((HandlerType::IntoAbi, path));
+                }
+                path if path.is_ident(FROM_ABI_WITH) => {
+                    let path = parse_path(meta)?;
+                    attributes.with_handlers.push((HandlerType::FromAbi, path));
+                }
+
+                path => {
+                    let path = path.get_ident().map(|x| x.to_string()).unwrap_or_default();
+                    return Err(Error::new(
+                        attr.span(),
+                        format!("Parameter {path} is not supported"),
+                    ));
+                }
             }
+
             Ok(())
-        });
+        })?;
     }
 
-    attributes
+    Ok(attributes)
 }
diff --git a/abi-proc/src/derive_from_abi.rs b/abi-proc/src/derive_from_abi.rs
@@ -1,18 +1,19 @@
 use crate::common;
-use crate::common::FieldAttributes;
+use crate::common::{FieldAttributes, HandlerType};
 use quote::quote;
+use syn::Error;
 
-pub fn impl_derive(input: syn::DeriveInput) -> Result<proc_macro2::TokenStream, syn::Error> {
+pub fn impl_derive(input: syn::DeriveInput) -> Result<proc_macro2::TokenStream, Error> {
     let data = match &input.data {
         syn::Data::Struct(data_struct) => data_struct,
         syn::Data::Enum(_) => {
-            return Err(syn::Error::new_spanned(
+            return Err(Error::new_spanned(
                 &input,
                 "FromAbi is not supported for enum",
             ))
         }
         syn::Data::Union(_) => {
-            return Err(syn::Error::new_spanned(
+            return Err(Error::new_spanned(
                 &input,
                 "FromAbi is not supported for unions",
             ))
@@ -28,8 +29,8 @@ pub fn impl_derive(input: syn::DeriveInput) -> Result<proc_macro2::TokenStream,
             continue;
         };
 
-        let attributes = common::extract_field_attributes(i.attrs.as_slice());
-        let token = construct_from_abi(name, &attributes);
+        let attributes = common::extract_field_attributes(i.attrs.as_slice())?;
+        let token = construct_from_abi(name, &attributes)?;
         struct_fields.push(token);
     }
 
@@ -51,19 +52,35 @@ pub fn impl_derive(input: syn::DeriveInput) -> Result<proc_macro2::TokenStream,
 pub fn construct_from_abi(
     field_name: &syn::Ident,
     attrs: &FieldAttributes,
-) -> proc_macro2::TokenStream {
-    match &attrs.custom_handler {
-        Some(handler) => {
-            quote!(#field_name: #handler::from_abi(
+) -> Result<proc_macro2::TokenStream, Error> {
+    attrs.check()?;
+
+    for (ty, path) in &attrs.with_handlers {
+        if !matches!(ty, HandlerType::FromAbi) {
+            continue;
+        }
+
+        return Ok(quote! {
+            #field_name: #path(
                 iter.next()
                 .ok_or(anyhow::anyhow!("unable to get field from abi"))?.value.clone())?
-            )
+        });
+    }
+
+    //fallback to mod handler if present
+    Ok(match &attrs.mod_handler {
+        Some(path) => {
+            quote! {
+                #field_name: #path::from_abi(
+                    iter.next()
+                    .ok_or(anyhow::anyhow!("unable to get field from abi"))?.value.clone())?
+            }
         }
         None => {
             quote!(#field_name: <_>::from_abi(
                 iter.next()
                 .ok_or(anyhow::anyhow!("unable to get field from abi"))?.value.clone())?
             )
         }
-    }
+    })
 }
diff --git a/abi-proc/src/derive_into_abi.rs b/abi-proc/src/derive_into_abi.rs
@@ -1,18 +1,19 @@
 use crate::common;
-use crate::common::FieldAttributes;
+use crate::common::{FieldAttributes, HandlerType};
 use quote::quote;
+use syn::Error;
 
-pub fn impl_derive(input: syn::DeriveInput) -> Result<proc_macro2::TokenStream, syn::Error> {
+pub fn impl_derive(input: syn::DeriveInput) -> Result<proc_macro2::TokenStream, Error> {
     let data = match input.data {
         syn::Data::Struct(data_struct) => data_struct,
         syn::Data::Enum(_) => {
-            return Err(syn::Error::new_spanned(
+            return Err(Error::new_spanned(
                 &input,
                 "IntoAbi is not supported for enum",
             ))
         }
         syn::Data::Union(_) => {
-            return Err(syn::Error::new_spanned(
+            return Err(Error::new_spanned(
                 &input,
                 "IntoAbi is not supported for unions",
             ))
@@ -28,8 +29,8 @@ pub fn impl_derive(input: syn::DeriveInput) -> Result<proc_macro2::TokenStream,
             continue;
         };
 
-        let attributes = common::extract_field_attributes(i.attrs.as_slice());
-        let token = construct_named_abi_value(name, &attributes);
+        let attributes = common::extract_field_attributes(i.attrs.as_slice())?;
+        let token = construct_named_abi_value(name, &attributes)?;
 
         tuple.push(token);
     }
@@ -61,15 +62,28 @@ pub fn impl_derive(input: syn::DeriveInput) -> Result<proc_macro2::TokenStream,
 pub fn construct_named_abi_value(
     field_name: &syn::Ident,
     attrs: &FieldAttributes,
-) -> proc_macro2::TokenStream {
+) -> Result<proc_macro2::TokenStream, Error> {
+    attrs.check()?;
+
     let to_change = match &attrs.custom_name {
         Some(custom) => custom.to_string(),
         None => field_name.to_string(),
     };
     let custom_name = to_change.to_string();
 
-    match &attrs.custom_handler {
+    for (ty, path) in &attrs.with_handlers {
+        if !matches!(ty, HandlerType::IntoAbi) {
+            continue;
+        }
+
+        return Ok(quote! {
+            quote!(#path(self.#field_name.clone()).named(#custom_name))
+        });
+    }
+
+    //fallback to mod handler if present
+    Ok(match &attrs.mod_handler {
         Some(handler) => quote!(#handler::into_abi(self.#field_name.clone()).named(#custom_name)),
         None => quote!(self.#field_name.clone().into_abi().named(#custom_name)),
-    }
+    })
 }
diff --git a/abi-proc/src/derive_with_abi_type.rs b/abi-proc/src/derive_with_abi_type.rs
@@ -1,18 +1,19 @@
 use crate::common;
-use crate::common::FieldAttributes;
+use crate::common::{FieldAttributes, HandlerType};
 use quote::quote;
+use syn::Error;
 
-pub fn impl_derive(input: syn::DeriveInput) -> Result<proc_macro2::TokenStream, syn::Error> {
+pub fn impl_derive(input: syn::DeriveInput) -> Result<proc_macro2::TokenStream, Error> {
     let data = match input.data {
         syn::Data::Struct(data_struct) => data_struct,
         syn::Data::Enum(_) => {
-            return Err(syn::Error::new_spanned(
+            return Err(Error::new_spanned(
                 &input,
                 "WithAbiType is not supported for enum",
             ))
         }
         syn::Data::Union(_) => {
-            return Err(syn::Error::new_spanned(
+            return Err(Error::new_spanned(
                 &input,
                 "WithAbiType is not supported for unions",
             ))
@@ -28,8 +29,8 @@ pub fn impl_derive(input: syn::DeriveInput) -> Result<proc_macro2::TokenStream,
             continue;
         };
 
-        let attributes = common::extract_field_attributes(i.attrs.as_slice());
-        let token = construct_with_abi_type(name, &i.ty, attributes);
+        let attributes = common::extract_field_attributes(i.attrs.as_slice())?;
+        let token = construct_with_abi_type(name, &i.ty, attributes)?;
 
         tuple.push(token);
     }
@@ -53,15 +54,28 @@ pub fn construct_with_abi_type(
     field_name: &syn::Ident,
     ty: &syn::Type,
     attrs: FieldAttributes,
-) -> proc_macro2::TokenStream {
+) -> Result<proc_macro2::TokenStream, Error> {
+    attrs.check()?;
+
     let to_change = match &attrs.custom_name {
         Some(custom) => custom.to_string(),
         None => field_name.to_string(),
     };
     let custom_name = to_change.to_string();
 
-    match &attrs.custom_handler {
+    for (ty, path) in &attrs.with_handlers {
+        if !matches!(ty, HandlerType::AbiType) {
+            continue;
+        }
+
+        return Ok(quote! {
+            #path().named(#custom_name)
+        });
+    }
+
+    //fallback to mod handler if present
+    Ok(match &attrs.mod_handler {
         Some(handler) => quote!(#handler::abi_type().named(#custom_name)),
         None => quote!(<#ty>::abi_type().named(#custom_name)),
-    }
+    })
 }
