From 9574da6f8715aa3ae4f5c999ab59d911528b381d Mon Sep 17 00:00:00 2001
From: Brendan Allan <brendonovich@outlook.com>
Date: Thu, 5 Oct 2023 13:00:38 +0800
Subject: [PATCH] make collect_commands a proc macro to support generic
 commands

---
 examples/app/src-tauri/src/main.rs |  7 +++-
 macros/src/collect_commands.rs     | 66 ++++++++++++++++++++++++++++++
 macros/src/lib.rs                  |  7 ++++
 src/lib.rs                         | 14 +------
 4 files changed, 80 insertions(+), 14 deletions(-)
 create mode 100644 macros/src/collect_commands.rs

diff --git a/examples/app/src-tauri/src/main.rs b/examples/app/src-tauri/src/main.rs
index 339e3bf..3f4258c 100644
--- a/examples/app/src-tauri/src/main.rs
+++ b/examples/app/src-tauri/src/main.rs
@@ -27,6 +27,10 @@ fn has_error() -> Result<&'static str, i32> {
     Err(32)
 }
 
+#[tauri::command]
+#[specta::specta]
+fn generic<T: tauri::Runtime>(app: tauri::AppHandle<T>) {}
+
 mod nested {
     use super::*;
 
@@ -57,7 +61,8 @@ fn main() {
                 hello_world,
                 goodbye_world,
                 has_error,
-                nested::some_struct
+                nested::some_struct,
+                generic::<tauri::Wry>
             ])
             .events(tauri_specta::collect_events![DemoEvent, EmptyEvent])
             .config(specta::ts::ExportConfig::default().formatter(specta::ts::prettier));
diff --git a/macros/src/collect_commands.rs b/macros/src/collect_commands.rs
new file mode 100644
index 0000000..194bbd0
--- /dev/null
+++ b/macros/src/collect_commands.rs
@@ -0,0 +1,66 @@
+use proc_macro2::{Ident, TokenStream};
+use quote::quote;
+use syn::{
+    bracketed,
+    parse::{Parse, ParseStream},
+    parse_macro_input,
+    punctuated::Punctuated,
+    ItemStruct, Path, Token,
+};
+
+pub struct Input {
+    type_map: Option<Ident>,
+    paths: Punctuated<Path, Token![,]>,
+}
+
+mod kw {
+    syn::custom_keyword!(type_map);
+}
+
+impl Parse for Input {
+    fn parse(input: ParseStream) -> syn::Result<Self> {
+        Ok(Self {
+            type_map: {
+                if input.peek(kw::type_map) && input.peek2(Token![:]) {
+                    input.parse::<kw::type_map>()?;
+                    input.parse::<Token![:]>()?;
+                    Some(input.parse()?)
+                } else {
+                    None
+                }
+            },
+            paths: Punctuated::parse_terminated(input)?,
+        })
+    }
+}
+
+pub fn proc_macro(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    let Input { type_map, paths } = parse_macro_input!(input as Input);
+
+    let tauri_paths = paths.iter().map(|p| {
+        let Path {
+            leading_colon,
+            segments,
+        } = p;
+
+        let segments = segments.iter().map(|s| &s.ident);
+
+        quote!(#leading_colon #(#segments)::*)
+    });
+
+    let type_map = type_map
+        .map(|i| quote!(#i))
+        .unwrap_or_else(|| quote!(::specta::TypeMap::new()));
+
+    let body = quote! {(
+        ::specta::collect_functions![type_map; #paths],
+        ::tauri::generate_handler![#(#tauri_paths),*],
+    )};
+
+    quote! {{
+        let mut type_map = #type_map;
+
+        #body
+    }}
+    .into()
+}
diff --git a/macros/src/lib.rs b/macros/src/lib.rs
index d7aacc0..e59e482 100644
--- a/macros/src/lib.rs
+++ b/macros/src/lib.rs
@@ -1,3 +1,5 @@
+mod collect_commands;
+
 use heck::ToKebabCase;
 use quote::quote;
 use syn::{parse_macro_input, DeriveInput};
@@ -17,3 +19,8 @@ pub fn derive_type(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     }
     .into()
 }
+
+#[proc_macro]
+pub fn collect_commands(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    collect_commands::proc_macro(input)
+}
diff --git a/src/lib.rs b/src/lib.rs
index 593928d..df1660f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -150,19 +150,7 @@ pub use event::*;
 pub type CollectCommandsTuple<TInvokeHandler> =
     (specta::functions::CollectFunctionsResult, TInvokeHandler);
 
-#[macro_export]
-macro_rules! collect_commands {
- 	(type_map: $type_map:ident, $($command:path),*) => {
-        (
-        	specta::collect_functions![$type_map; $($command),*],
-       		::tauri::generate_handler![$($command),*],
-        )
-    };
-    ($($command:path),*) => {{
-        let mut type_map = specta::TypeMap::default();
-        $crate::collect_commands![type_map: type_map, $($command),*]
-    }};
-}
+pub use tauri_specta_macros::collect_commands;
 
 // TODO
 // #[cfg(doctest)]