diff --git a/src/lib.rs b/src/lib.rs index 7869c89..b542010 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,17 +9,32 @@ pub fn id_derive(input: TokenStream) -> TokenStream { // Parse the input tokens into a syntax tree let ast = parse_macro_input!(input as DeriveInput); let name = &ast.ident; + let ty = match &ast.data { + syn::Data::Struct(s) => { + if let syn::Fields::Unnamed(fields) = &s.fields { + if fields.unnamed.len() != 1 { + return syn::Error::new_spanned(&ast, "Expected exactly one field.").to_compile_error().into(); + } + &fields.unnamed[0].ty + } else { + return syn::Error::new_spanned(&ast, "Expected a tuple struct.").to_compile_error().into(); + } + } + _ => { + return syn::Error::new_spanned(&ast, "Expected a struct.").to_compile_error().into(); + } + }; // Generate the code for the implementation let expanded = quote! { impl #name { #[inline] - pub const fn new(id: u32) -> Self { + pub const fn new(id: #ty) -> Self { Self(id) } #[inline] - pub const fn get(self) -> u32 { + pub const fn get(self) -> #ty { self.0 } }