diff options
Diffstat (limited to 'src/token_stream.rs')
| -rw-r--r-- | src/token_stream.rs | 260 |
1 files changed, 52 insertions, 208 deletions
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"))), ) }); |