refactor: simplify replacing identifiers

2 files changed, 62 insertions, 237 deletions

diff --git a/src/repeat.rs b/src/repeat.rs
index 2f55ee8..5d9ece7 100644
--- a/src/repeat.rs
+++ b/src/repeat.rs
@@ -19,18 +19,12 @@ pub fn for_each_opengl_command_impl(
     opengl_registry: &Registry,
 ) -> TokenStream
 {
-    let gl_func_ident_index_paths = input_stream
-        .find_all_ident(&Ident::new(OPENGL_CMD_IDENT_REPLACE, Span::call_site()));
-
-    let gl_func_ret_type_index_paths = input_stream
-        .find_all_ident(&Ident::new(OPENGL_CMD_RET_TYPE_REPLACE, Span::call_site()));
-
     let out: TokenStream = opengl_registry
         .commands()
         .iter()
         .map(|gl_command| {
-            let stream = input_stream.replace_tokens(
-                &gl_func_ident_index_paths,
+            let stream = input_stream.replace_ident(
+                &Ident::new(OPENGL_CMD_IDENT_REPLACE, Span::call_site()),
                 &TokenTree::Ident(Ident::new(
                     gl_command.prototype().name(),
                     Span::call_site(),
@@ -39,19 +33,16 @@ pub fn for_each_opengl_command_impl(
 
             let return_type = gl_command.prototype().return_type();
 
-            let stream = stream.replace_tokens(
-                &gl_func_ret_type_index_paths,
+            let stream = stream.replace_ident(
+                &Ident::new(OPENGL_CMD_RET_TYPE_REPLACE, Span::call_site()),
                 &TokenTree::Group(Group::new(
                     Delimiter::None,
                     TokenStream::from_str(&fix_type(return_type)).unwrap(),
                 )),
             );
 
-            let gl_func_args_paths = stream
-                .find_all_ident(&Ident::new(OPENGL_CMD_ARGS_REPLACE, Span::call_site()));
-
-            let stream = stream.replace_tokens(
-                &gl_func_args_paths,
+            let stream = stream.replace_ident(
+                &Ident::new(OPENGL_CMD_ARGS_REPLACE, Span::call_site()),
                 &TokenTree::Group(Group::new(
                     Delimiter::None,
                     TokenStream::from_str(
@@ -71,13 +62,8 @@ pub fn for_each_opengl_command_impl(
                 )),
             );
 
-            let gl_func_arg_names_paths = stream.find_all_ident(&Ident::new(
-                OPENGL_CMD_ARG_NAMES_REPLACE,
-                Span::call_site(),
-            ));
-
-            let stream = stream.replace_tokens(
-                &gl_func_arg_names_paths,
+            let stream = stream.replace_ident(
+                &Ident::new(OPENGL_CMD_ARG_NAMES_REPLACE, Span::call_site()),
                 &TokenTree::Group(Group::new(
                     Delimiter::None,
                     TokenStream::from_str(
@@ -92,13 +78,8 @@ pub fn for_each_opengl_command_impl(
                 )),
             );
 
-            let gl_func_arg_types_paths = stream.find_all_ident(&Ident::new(
-                OPENGL_CMD_ARG_TYPES_REPLACE,
-                Span::call_site(),
-            ));
-
-            stream.replace_tokens(
-                &gl_func_arg_types_paths,
+            stream.replace_ident(
+                &Ident::new(OPENGL_CMD_ARG_TYPES_REPLACE, Span::call_site()),
                 &TokenTree::Group(Group::new(
                     Delimiter::None,
                     TokenStream::from_str(
diff --git a/src/token_stream.rs b/src/token_stream.rs
index 5b755c9..5897529 100644
--- a/src/token_stream.rs
+++ b/src/token_stream.rs
@@ -1,165 +1,44 @@
 use proc_macro2::{Group, Ident, TokenStream, TokenTree};
 
-pub use self::index_path::IndexPath;
-
 #[allow(clippy::module_name_repetitions)]
 pub trait TokenStreamExt
 {
-    fn find_all_ident(&self, target_ident: &Ident) -> Vec<IndexPath>;
-
-    fn replace_tokens(
+    /// Recursively replaces all occurences of a identifier with a token tree.
+    fn replace_ident(
         &self,
-        index_paths_to_replace: &[IndexPath],
+        target_ident: &Ident,
         substitution: &TokenTree,
     ) -> TokenStream;
 }
 
 impl TokenStreamExt for TokenStream
 {
-    fn find_all_ident(&self, target_ident: &Ident) -> Vec<IndexPath>
-    {
-        let mut found_indices = Vec::new();
-
-        recurse_find_all_ident(
-            self.clone(),
-            target_ident,
-            &mut found_indices,
-            &IndexPath::new(),
-        );
-
-        found_indices
-    }
-
-    fn replace_tokens(
-        &self,
-        index_paths_to_replace: &[IndexPath],
-        substitution: &TokenTree,
-    ) -> TokenStream
+    fn replace_ident(&self, target_ident: &Ident, substitution: &TokenTree)
+        -> TokenStream
     {
-        self.clone()
-            .into_iter()
-            .enumerate()
-            .map(|(index, mut token_tree)| {
-                for index_path in index_paths_to_replace
-                    .iter()
-                    .filter(|path| path.indices()[0] == index)
-                {
-                    token_tree = match token_tree {
-                        TokenTree::Ident(_) => substitution.clone(),
-                        TokenTree::Group(group) => TokenTree::Group(Group::new(
-                            group.delimiter(),
-                            recurse_replace_tokens(
-                                group.stream(),
-                                substitution,
-                                &index_path.indices()[1..],
-                            ),
-                        )),
-                        tt => tt,
-                    }
-                }
-
-                token_tree
-            })
-            .collect()
+        recurse_replace_ident(self.clone(), target_ident, substitution)
     }
 }
 
-fn recurse_find_all_ident(
-    input_stream: TokenStream,
-    target_ident: &Ident,
-    found_indices: &mut Vec<IndexPath>,
-    current_index_path: &IndexPath,
-)
-{
-    for (index, token_tree) in input_stream.into_iter().enumerate() {
-        match token_tree {
-            TokenTree::Ident(ident) if &ident == target_ident => {
-                let mut index_path = current_index_path.clone();
-
-                index_path.push_index(index);
-
-                found_indices.push(index_path);
-            }
-            TokenTree::Group(group) => {
-                let mut index_path = current_index_path.clone();
-
-                index_path.push_index(index);
-
-                recurse_find_all_ident(
-                    group.stream(),
-                    target_ident,
-                    found_indices,
-                    &index_path,
-                );
-            }
-            _ => {}
-        }
-    }
-}
-
-fn recurse_replace_tokens(
+fn recurse_replace_ident(
     token_stream: TokenStream,
+    target_ident: &Ident,
     substitution: &TokenTree,
-    indices: &[usize],
 ) -> TokenStream
 {
     token_stream
         .into_iter()
-        .enumerate()
-        .map(|(index, token_tree)| match token_tree {
-            TokenTree::Ident(_) if index == indices[0] => substitution.clone(),
-            TokenTree::Group(group) if index == indices[0] => {
-                TokenTree::Group(Group::new(
-                    group.delimiter(),
-                    recurse_replace_tokens(group.stream(), substitution, &indices[1..]),
-                ))
-            }
+        .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,
         })
         .collect()
 }
 
-mod index_path
-{
-    #[derive(Debug, Clone, PartialEq, Eq)]
-    pub struct IndexPath
-    {
-        indices: Vec<usize>,
-    }
-
-    impl IndexPath
-    {
-        pub fn new() -> Self
-        {
-            Self {
-                indices: Vec::new(),
-            }
-        }
-
-        pub fn push_index(&mut self, index: usize)
-        {
-            self.indices.push(index);
-        }
-
-        pub fn indices(&self) -> &[usize]
-        {
-            &self.indices
-        }
-    }
-
-    impl<IntoIter> From<IntoIter> for IndexPath
-    where
-        IntoIter: IntoIterator<Item = usize>,
-    {
-        fn from(value: IntoIter) -> Self
-        {
-            Self {
-                indices: value.into_iter().collect(),
-            }
-        }
-    }
-}
-
 #[cfg(test)]
 mod tests
 {
@@ -174,59 +53,15 @@ mod tests
     use super::*;
 
     #[test]
-    fn find_all_ident_works()
+    fn replace_ident_works()
     {
         assert_eq!(
             quote! {
                 let abc = xyz;
             }
-            .find_all_ident(&Ident::new("xyz", Span::call_site())),
-            vec![IndexPath::from([3])]
-        );
-
-        assert_eq!(
-            quote! {
-                let abc = (xyz, "123");
-            }
-            .find_all_ident(&Ident::new("xyz", Span::call_site())),
-            vec![IndexPath::from([3, 0])]
-        );
-
-        assert_eq!(
-            quote! {
-                return ("123", (yo, 180, xyz));
-            }
-            .find_all_ident(&Ident::new("xyz", Span::call_site())),
-            vec![IndexPath::from([1, 2, 4])]
-        );
-    }
-
-    #[test]
-    fn find_all_ident_works_with_multiple()
-    {
-        assert_eq!(
-            quote! {
-                unsafe {
-                    functions::xyz = FunctionPtr::new_initialized(
-                        get_proc_addr(stringify!(xyz))
-                    );
-                }
-            }
-            .find_all_ident(&Ident::new("xyz", Span::call_site()),),
-            vec![IndexPath::from([1, 3]), IndexPath::from([1, 9, 1, 2, 0])]
-        );
-    }
-
-    #[test]
-    fn recurse_replace_tokens_works()
-    {
-        assert_eq!(
-            recurse_replace_tokens(
-                quote! {
-                    let abc = xyz;
-                },
+            .replace_ident(
+                &format_ident!("xyz"),
                 &TokenTree::Ident(Ident::new("foo", Span::call_site())),
-                &[3]
             )
             .to_string(),
             quote! {
@@ -236,12 +71,12 @@ mod tests
         );
 
         assert_eq!(
-            recurse_replace_tokens(
-                quote! {
-                    let abc = (xyz, "123");
-                },
+            quote! {
+                let abc = (xyz, "123");
+            }
+            .replace_ident(
+                &format_ident!("xyz"),
                 &TokenTree::Ident(Ident::new("foo", Span::call_site())),
-                &[3, 0]
             )
             .to_string(),
             quote! {
@@ -251,12 +86,12 @@ mod tests
         );
 
         assert_eq!(
-            recurse_replace_tokens(
-                quote! {
-                    let abc = (hello, "123").iter_map(|_| xyz);
-                },
+            quote! {
+                let abc = (hello, "123").iter_map(|_| xyz);
+            }
+            .replace_ident(
+                &format_ident!("xyz"),
                 &TokenTree::Ident(Ident::new("foo", Span::call_site())),
-                &[6, 3]
             )
             .to_string(),
             quote! {
@@ -264,28 +99,37 @@ mod tests
             }
             .to_string()
         );
-    }
 
-    #[bench]
-    fn bench_find_all_ident(bencher: &mut Bencher)
-    {
-        let input_strem = quote! {
-            let foo = |abc| {
-                let x = "foo";
+        assert_eq!(
+            quote! {
+                fn xyz(bar: &Bar) {
+                    let foo = "baz";
 
-                let y = a_foo;
+                    foo
+                }
 
-                let foo = "Hello";
+                println!("Hello {}", xyz());
+            }
+            .replace_ident(
+                &format_ident!("xyz"),
+                &TokenTree::Ident(Ident::new("foo", Span::call_site())),
+            )
+            .to_string(),
+            quote! {
+                fn foo(bar: &Bar) {
+                    let foo = "baz";
 
-                foo.to_string()
-            };
-        };
+                    foo
+                }
 
-        bencher.iter(|| input_strem.find_all_ident(&black_box(format_ident!("foo"))));
+                println!("Hello {}", foo());
+            }
+            .to_string()
+        );
     }
 
     #[bench]
-    fn bench_replace_tokens(bencher: &mut Bencher)
+    fn bench_replace_ident(bencher: &mut Bencher)
     {
         let input_strem = quote! {
             let foo = |abc| {
@@ -300,8 +144,8 @@ mod tests
         };
 
         bencher.iter(|| {
-            input_strem.replace_tokens(
-                black_box(&[[1].into(), [6, 11].into(), [6, 15].into()]),
+            input_strem.replace_ident(
+                black_box(&format_ident!("foo")),
                 black_box(&TokenTree::Ident(format_ident!("foobar"))),
             )
         });