feat(util-macros): allow passing extra args to sub macro callbacks

1 files changed, 44 insertions, 3 deletions

diff --git a/util-macros/src/lib.rs b/util-macros/src/lib.rs
index 036ffc8..26f1936 100644
--- a/util-macros/src/lib.rs
+++ b/util-macros/src/lib.rs
@@ -1,10 +1,12 @@
 use proc_macro::{TokenStream, TokenTree};
 use quote::quote;
 
-/// Subtracts two numbers and calls a given callback macro with the result.
+/// Subtracts two numbers and calls a given callback macro with the result. Optionally, a
+/// additional argument (delimeted) can be given which will also be passed to the
+/// callback.
 ///
 /// # Input
-/// `number - number, callback`
+/// `$num_a - $num_b, $callback $(, $user_data)?`
 ///
 /// # Examples
 /// ```
@@ -23,6 +25,23 @@ use quote::quote;
 /// ```
 /// <br>
 ///
+/// The callback macro can be called with extra arguments.
+/// ```
+/// # use std::any::TypeId;
+/// use util_macros::sub;
+///
+/// macro_rules! sub_cb {
+///     ($num: literal, $to_multiply: literal) => {
+///         $num * $to_multiply
+///     };
+/// }
+///
+/// type Foo = [u8; sub!(5 - 2, sub_cb, (20))];
+///
+/// assert_eq!(TypeId::of::<Foo>(), TypeId::of::<[u8; 60]>());
+/// ```
+/// <br>
+///
 /// The callback is called with the identifier `overflow` if a overflow occurs.
 /// ```
 /// # use std::any::TypeId;
@@ -84,6 +103,28 @@ pub fn sub(input: TokenStream) -> TokenStream
         }
     };
 
+    let opt_user_data = input_tt_iter
+        .next()
+        .map(|comma_tt| {
+            match comma_tt {
+                TokenTree::Punct(punct) if punct.as_char() == ',' => {}
+                _ => {
+                    panic!("Expected a ',' token");
+                }
+            };
+
+            let user_data_tt = input_tt_iter.next().unwrap();
+
+            let TokenTree::Group(group) = user_data_tt else {
+                panic!("User data must be a delimeted")
+            };
+
+            let inside: proc_macro2::TokenStream = group.stream().into();
+
+            quote! {, #inside }
+        })
+        .unwrap_or_default();
+
     let Some(subtracted) = num_a.checked_sub(num_b) else {
         return quote! {
             #cb_ident!(overflow)
@@ -94,7 +135,7 @@ pub fn sub(input: TokenStream) -> TokenStream
     let subtracted_lit = proc_macro2::Literal::u32_unsuffixed(subtracted);
 
     quote! {
-        #cb_ident!(#subtracted_lit)
+        #cb_ident!(#subtracted_lit #opt_user_data)
     }
     .into()
 }