refactor: allow conversion method to wrap type

1 files changed, 23 insertions, 21 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 5385181..4512141 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -203,11 +203,11 @@ pub fn dynamize(_attr: TokenStream, input: TokenStream) -> TokenStream {
                                 parsed_method.type_param_transforms.get(type_ident)
                             {
                                 let args = (0..transforms.len()).map(|i| format_ident!("a{}", i));
-                                let mut calls: Vec<_> =
-                                    args.clone().map(|i| i.into_token_stream()).collect();
-                                for i in 0..calls.len() {
-                                    transforms[i].append_conversion(&mut calls[i]);
-                                }
+                                let calls: Vec<_> = args
+                                    .clone()
+                                    .enumerate()
+                                    .map(|(idx, i)| transforms[idx].convert(i.into_token_stream()))
+                                    .collect();
                                 let move_opt = new_method.sig.asyncness.map(|_| quote! {move});
                                 quote!(#move_opt |#(#args),*| #ident(#(#calls),*))
                             } else {
@@ -275,7 +275,7 @@ pub fn dynamize(_attr: TokenStream, input: TokenStream) -> TokenStream {
         if new_method.sig.asyncness.is_some() {
             expr.extend(quote! {.await})
         }
-        parsed_method.return_type.append_conversion(&mut expr);
+        let expr = parsed_method.return_type.convert(expr);
 
         method_impls.push(ImplItemMethod {
             attrs: Vec::new(),
@@ -362,27 +362,29 @@ where
 }
 
 impl TypeTransform {
-    fn append_conversion(&self, stream: &mut proc_macro2::TokenStream) {
+    fn convert(&self, arg: proc_macro2::TokenStream) -> proc_macro2::TokenStream {
         match self {
-            TypeTransform::Into => stream.extend(quote! {.into()}),
+            TypeTransform::Into => quote! {#arg.into()},
             TypeTransform::Map(opt) => {
-                let mut inner = quote!(x);
-                opt.append_conversion(&mut inner);
-                stream.extend(quote! {.map(|x| #inner)})
+                let inner = opt.convert(quote!(x));
+                quote! {#arg.map(|x| #inner)}
             }
             TypeTransform::Result(ok, err) => {
-                if !matches!(ok.as_ref(), TypeTransform::NoOp) {
-                    let mut inner = quote!(x);
-                    ok.append_conversion(&mut inner);
-                    stream.extend(quote! {.map(|x| #inner)})
-                }
-                if !matches!(err.as_ref(), TypeTransform::NoOp) {
-                    let mut inner = quote!(x);
-                    err.append_conversion(&mut inner);
-                    stream.extend(quote! {.map_err(|x| #inner)})
+                let map_ok = !matches!(ok.as_ref(), TypeTransform::NoOp);
+                let map_err = !matches!(err.as_ref(), TypeTransform::NoOp);
+                if map_ok && map_err {
+                    let ok_inner = ok.convert(quote!(x));
+                    let err_inner = err.convert(quote!(x));
+                    quote! {#arg.map(|x| #ok_inner).map_err(|x| #err_inner)}
+                } else if map_ok {
+                    let ok_inner = ok.convert(quote!(x));
+                    quote! {#arg.map(|x| #ok_inner)}
+                } else {
+                    let err_inner = err.convert(quote!(x));
+                    quote! {#arg.map_err(|x| #err_inner)}
                 }
             }
-            _other => {}
+            _other => arg,
         }
     }
 }