refactor: improve replacement syntax

1 files changed, 60 insertions, 17 deletions

diff --git a/src/token_stream.rs b/src/token_stream.rs
index 5897529..1bbf431 100644
--- a/src/token_stream.rs
+++ b/src/token_stream.rs
@@ -1,5 +1,7 @@
 use proc_macro2::{Group, Ident, TokenStream, TokenTree};
 
+use crate::iter::IteratorExt;
+
 #[allow(clippy::module_name_repetitions)]
 pub trait TokenStreamExt
 {
@@ -16,7 +18,7 @@ impl TokenStreamExt for TokenStream
     fn replace_ident(&self, target_ident: &Ident, substitution: &TokenTree)
         -> TokenStream
     {
-        recurse_replace_ident(self.clone(), target_ident, substitution)
+        recurse_replace_ident(self.clone(), target_ident, substitution, '#')
     }
 }
 
@@ -24,17 +26,50 @@ fn recurse_replace_ident(
     token_stream: TokenStream,
     target_ident: &Ident,
     substitution: &TokenTree,
+    prefix: char,
 ) -> TokenStream
 {
+    let mut prev_token_was_prefix = false;
+
     token_stream
         .into_iter()
-        .map(|token_tree| match token_tree {
-            TokenTree::Ident(ident) if &ident == target_ident => substitution.clone(),
-            TokenTree::Group(group) => TokenTree::Group(Group::new(
-                group.delimiter(),
-                recurse_replace_ident(group.stream(), target_ident, substitution),
-            )),
-            tt => tt,
+        .windows()
+        .filter_map(|(token_tree, next_token_tree)| match token_tree {
+            TokenTree::Punct(punct) if punct.as_char() == prefix => {
+                if matches!(next_token_tree, Some(TokenTree::Ident(ident)) if &ident == target_ident) {
+                    prev_token_was_prefix = true;
+
+                    None
+                }
+                else {
+                    prev_token_was_prefix = false;
+
+                    Some(TokenTree::Punct(punct))
+                }
+            }
+            TokenTree::Ident(ident) if prev_token_was_prefix && &ident == target_ident => {
+                prev_token_was_prefix = false;
+
+                Some(substitution.clone())
+            }
+            TokenTree::Ident(_) if prev_token_was_prefix => {
+                prev_token_was_prefix = false;
+
+                Some(substitution.clone())
+            }
+            TokenTree::Group(group) => {
+                prev_token_was_prefix = false;
+
+                Some(TokenTree::Group(Group::new(
+                    group.delimiter(),
+                    recurse_replace_ident(group.stream(), target_ident, substitution, prefix),
+                )))
+            }
+            tt => {
+                prev_token_was_prefix = false;
+
+                Some(tt)
+            }
         })
         .collect()
 }
@@ -46,7 +81,7 @@ mod tests
 
     use std::hint::black_box;
 
-    use proc_macro2::Span;
+    use proc_macro2::{Punct, Spacing, Span};
     use quote::{format_ident, quote};
     use test::Bencher;
 
@@ -55,9 +90,11 @@ mod tests
     #[test]
     fn replace_ident_works()
     {
+        let ht = Punct::new('#', Spacing::Alone);
+
         assert_eq!(
             quote! {
-                let abc = xyz;
+                let abc = #ht xyz;
             }
             .replace_ident(
                 &format_ident!("xyz"),
@@ -72,7 +109,7 @@ mod tests
 
         assert_eq!(
             quote! {
-                let abc = (xyz, "123");
+                let abc = (#ht xyz, "123");
             }
             .replace_ident(
                 &format_ident!("xyz"),
@@ -87,7 +124,7 @@ mod tests
 
         assert_eq!(
             quote! {
-                let abc = (hello, "123").iter_map(|_| xyz);
+                let abc = (hello, "123").iter_map(|_| #ht xyz);
             }
             .replace_ident(
                 &format_ident!("xyz"),
@@ -102,13 +139,15 @@ mod tests
 
         assert_eq!(
             quote! {
-                fn xyz(bar: &Bar) {
+                fn #ht xyz(bar: &Bar) {
                     let foo = "baz";
 
                     foo
                 }
 
-                println!("Hello {}", xyz());
+                let xyz = "123";
+
+                println!("Hello {}", #ht xyz());
             }
             .replace_ident(
                 &format_ident!("xyz"),
@@ -122,6 +161,8 @@ mod tests
                     foo
                 }
 
+                let xyz = "123";
+
                 println!("Hello {}", foo());
             }
             .to_string()
@@ -131,15 +172,17 @@ mod tests
     #[bench]
     fn bench_replace_ident(bencher: &mut Bencher)
     {
+        let ht = Punct::new('#', Spacing::Alone);
+
         let input_strem = quote! {
-            let foo = |abc| {
+            let #ht foo = |abc| {
                 let x = "foo";
 
                 let y = a_foo;
 
-                let foo = "Hello";
+                let #ht foo = "Hello";
 
-                foo.to_string()
+                #ht foo.to_string()
             };
         };