add type_check sol attr

alloy-rs · Dec 30, 2023 · dbff597 · dbff597
1 parent bb879ec
commit dbff597
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 2 deletions.
diff --git a/crates/sol-macro/src/attr.rs b/crates/sol-macro/src/attr.rs
@@ -78,6 +78,8 @@ pub struct SolAttrs {
 
     pub bytecode: Option<LitStr>,
     pub deployed_bytecode: Option<LitStr>,
+
+    pub type_check: Option<LitStr>,
 }
 
 impl SolAttrs {
@@ -146,6 +148,8 @@ impl SolAttrs {
 
                     bytecode => bytes()?,
                     deployed_bytecode => bytes()?,
+
+                    type_check => lit()?,
                 };
                 Ok(())
             })?;
@@ -322,5 +326,10 @@ mod tests {
             #[sol(bytecode = "xyz")] => Err("expected hex literal"),
             #[sol(bytecode = "123")] => Err("expected even number of hex digits"),
         }
+
+        type_check {
+            #[sol(type_check = "my_function")] => Ok(sol_attrs! {type_check: parse_quote!("my_function")} ),
+            #[sol(type_check = "my_function1")] #[sol(type_check = "my_function2")] => Err(DUPLICATE_ERROR),
+        }
     }
 }
diff --git a/crates/sol-macro/src/expand/udt.rs b/crates/sol-macro/src/expand/udt.rs
@@ -10,9 +10,16 @@ use syn::Result;
 pub(super) fn expand(cx: &ExpCtxt<'_>, udt: &ItemUdt) -> Result<TokenStream> {
     let ItemUdt { name, ty, attrs, .. } = udt;
 
-    let (_sol_attrs, mut attrs) = crate::attr::SolAttrs::parse(attrs)?;
+    let (sol_attrs, mut attrs) = crate::attr::SolAttrs::parse(attrs)?;
     cx.type_derives(&mut attrs, Some(ty), true);
 
+    let type_check = if let Some(lit_str) = sol_attrs.type_check {
+        let func_path: syn::Path = syn::parse_str(&lit_str.value()).unwrap();
+        quote! { #func_path }
+    } else {
+        quote! { ::alloy_sol_types::private::just_ok }
+    };
+
     let underlying = expand_type(ty);
     let tokens = quote! {
         #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
@@ -90,7 +97,7 @@ pub(super) fn expand(cx: &ExpCtxt<'_>, udt: &ItemUdt) -> Result<TokenStream> {
             #[inline]
             fn type_check(token: &Self::Token<'_>) -> ::alloy_sol_types::Result<()> {
                 <#underlying as ::alloy_sol_types::SolType>::type_check(token)?;
-                ::alloy_sol_types::private::just_ok(token)
+                #type_check(token)
             }
 
             #[inline]

diff --git a/crates/sol-macro/src/lib.rs b/crates/sol-macro/src/lib.rs
@@ -117,6 +117,8 @@ mod json;
 ///   will emit a `static` item with the specified bytes.
 /// - `deployed_bytecode = <hex string literal>`: specifies the deployed bytecode of a contract.
 ///   This will emit a `static` item with the specified bytes.
+/// - `type_check = <string literal>`: specifies a function to be used to check an User Defined
+///   Type.
 ///
 /// ### Structs and enums
 ///