Fix oversight in new function handling

Bro, these macros killed me to write.
specta-rs · Jul 30, 2024 · fd0cfca · fd0cfca
1 parent d6c3b39
commit fd0cfca
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 7 deletions.
diff --git a/specta-macros/src/specta.rs b/specta-macros/src/specta.rs
@@ -96,8 +96,8 @@ pub fn attribute(item: proc_macro::TokenStream) -> syn::Result<proc_macro::Token
         #maybe_macro_export
         #[doc(hidden)]
         macro_rules! #wrapper {
-            // We take in `$function` from the invocation so we have `name::<concrete_generics_types>`
-            (@export_fn; $function:ident) => {{
+            // We take in `$function` from the invocation so we have `fn_name::<concrete_generics_types>`
+            (@export_fn; $function:path) => {{
                 fn export(type_map: &mut #crate_ref::TypeMap) -> #crate_ref::datatype::Function {
                     #crate_ref::internal::get_fn_datatype(
                         $function as fn(#(#arg_signatures),*) -> _,

diff --git a/specta/src/function/mod.rs b/specta/src/function/mod.rs
@@ -33,11 +33,27 @@ pub(crate) use specta_fn::SpectaFn;
 ///     assert_eq!(typ.result, Some(DataType::Primitive(PrimitiveType::bool)));
 /// }
 /// ```
+///
+// TODO: Hide implementation details in inner macro like Serde does
 #[doc(hidden)]
 #[macro_export]
 macro_rules! _fn_datatype {
-    ($function:path) => {
-        $crate::internal::paste! { [<__specta__fn__ $function>]!(@export_fn; $function) }
+    ([$($path:tt)*] [$($full:tt)*] [$last:tt]) => {
+        $crate::internal::paste! {
+            $($path)* [<__specta__fn__ $last>]!(@export_fn; $($full)*)
+        }
+    };
+    ([$($path:tt)*] [$($full:tt)*] [$($last:tt)?] $t:tt :: <$($g:path)*> $($rest:tt)*) => {
+        $crate::function::fn_datatype!([$($path)* $($last)*] [$($full)* $t::<$($g)*>] [$t] $($rest)*)
+    };
+    ([$($path:tt)*] [$($full:tt)*] [$($last:tt)?] $t:tt $($rest:tt)*) => {
+        $crate::function::fn_datatype!([$($path)* $($last)*] [$($full)* $t] [$t] $($rest)*)
+    };
+    () => {{
+            compile_error!("fn_datatype must be provided a function path as an argument");
+    }};
+    ($($rest:tt)*) => {
+        $crate::function::fn_datatype!([] [] [] $($rest)*)
     };
 }
 
@@ -60,16 +76,29 @@ macro_rules! _fn_datatype {
 ///     let functions = function::collect_functions![some_function](&mut TypeMap::default());
 /// }
 /// ````
+// TODO: Hide implementation details in inner macro like Serde does
 #[doc(hidden)]
 #[macro_export]
 macro_rules! _collect_functions {
-    ($($command:path),* $(,)?) => {{
-        fn export(type_map: &mut $crate::TypeMap) -> Vec<$crate::datatype::Function> {
-            vec![$($crate::fn_datatype!($command)(type_map)),*]
+    ($tm:ident [] [$($result:expr)*]) => {{
+        fn export($tm: &mut $crate::TypeMap) -> Vec<$crate::datatype::Function> {
+            vec![$($result),*]
         }
 
         export
     }};
+    ($tm:ident [$($parts:tt)*] [$($result:expr)*]) => {
+        $crate::function::collect_functions!($tm [] [$($result)* $crate::function::fn_datatype!($($parts)*)($tm)])
+    };
+    ($tm:ident [$($parts:tt)*] [$($result:expr)*] , $($rest:tt)*) => {
+        $crate::function::collect_functions!($tm [] [$($result)* $crate::function::fn_datatype!($($parts)*)($tm)] $($rest)*)
+    };
+    ($tm:ident [$($parts:tt)*] [$($result:expr)*] $t:tt $($rest:tt)*) => {
+        $crate::function::collect_functions!($tm [$($parts)* $t] [$($result)*] $($rest)*)
+    };
+    ($($command:tt)*) => {
+        $crate::function::collect_functions!(type_map [] [] $($command)*)
+    };
 }
 
 #[doc(inline)]