From 6bf6968ce62a82ed208e1cbadcecd42f079687f3 Mon Sep 17 00:00:00 2001
From: Omer Ozarslan <omerfaruko@gmail.com>
Date: Fri, 12 Apr 2024 18:25:32 -0700
Subject: [PATCH] Accept mut prefix before argument names

---
 src/lib.rs | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 174ab5e..f2c7dbb 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -44,16 +44,18 @@ impl ToTokens for FnArgName {
 /// Function argument with name and type.
 #[derive(Debug, Eq, PartialEq, Clone)]
 struct FnArg {
+    r#mut: Option<Token![mut]>,
     name: FnArgName,
     ty: Type,
 }
 
 impl Parse for FnArg {
     fn parse(input: ParseStream) -> Result<Self> {
+        let r#mut = input.parse()?;
         let name = input.parse()?;
         let _ = input.parse::<Token![:]>()?;
         let ty = input.parse()?;
-        Ok(Self { name, ty })
+        Ok(Self { r#mut, name, ty })
     }
 }
 
@@ -185,6 +187,7 @@ fn generate_trait_implementation(
     trait_name: &Ident,
     generic_params: Option<&Punctuated<GenericParam, Token![,]>>,
     FnArg {
+        r#mut: input_expr_mut,
         name: input_expr_name,
         ty: input_expr_type,
     }: &FnArg,
@@ -195,7 +198,7 @@ fn generate_trait_implementation(
     let generics = generic_params.map(|g| quote! {<#g>});
     quote! {
         impl #generics #trait_name<#input_expr_type> for #input_expr_type {
-            #default fn dispatch(#input_expr_name: #input_expr_type #(, #extra_args)*) -> #return_type {
+            #default fn dispatch(#input_expr_mut #input_expr_name: #input_expr_type #(, #extra_args)*) -> #return_type {
                 #body
             }
         }
@@ -407,4 +410,17 @@ mod tests {
             }
         );
     }
+
+    #[test]
+    fn parse_mut_arg() {
+        let arg: FnArg = parse_quote!(mut v: u8);
+        assert_eq!(
+            arg,
+            FnArg {
+                r#mut: Some(parse_quote!(mut)),
+                ty: parse_quote!(u8),
+                name: FnArgName::Ident(parse_quote!(v)),
+            }
+        );
+    }
 }