telety_impl/
find_and_replace.rs1use proc_macro2::{Group, Literal, Punct, TokenStream, TokenTree};
3use quote::{ToTokens, TokenStreamExt};
4use syn::{
5 parse::{Parse, ParseStream},
6 Ident,
7};
8
9pub enum SingleToken {
10 Ident(Ident),
11 Punct(Punct),
12 Literal(Literal),
13}
14
15impl From<Ident> for SingleToken {
16 fn from(value: Ident) -> Self {
17 Self::Ident(value)
18 }
19}
20
21impl From<Punct> for SingleToken {
22 fn from(value: Punct) -> Self {
23 Self::Punct(value)
24 }
25}
26
27impl From<Literal> for SingleToken {
28 fn from(value: Literal) -> Self {
29 Self::Literal(value)
30 }
31}
32
33impl Parse for SingleToken {
34 fn parse(input: ParseStream) -> syn::Result<Self> {
35 let tt: TokenTree = input.parse()?;
36 match tt {
37 TokenTree::Group(g) => Err(syn::Error::new(
38 g.span(),
39 "Only single tokens are allowed as needles",
40 )),
41 TokenTree::Ident(i) => Ok(Self::Ident(i)),
42 TokenTree::Punct(p) => Ok(Self::Punct(p)),
43 TokenTree::Literal(l) => Ok(Self::Literal(l)),
44 }
45 }
46}
47
48impl ToTokens for SingleToken {
49 fn to_tokens(&self, tokens: &mut TokenStream) {
50 match self {
51 SingleToken::Ident(i) => i.to_tokens(tokens),
52 SingleToken::Punct(p) => p.to_tokens(tokens),
53 SingleToken::Literal(l) => l.to_tokens(tokens),
54 }
55 }
56}
57
58impl PartialEq<TokenTree> for SingleToken {
59 fn eq(&self, other: &TokenTree) -> bool {
60 match (self, other) {
61 (SingleToken::Ident(a), TokenTree::Ident(b)) if a == b => true,
62 (SingleToken::Punct(a), TokenTree::Punct(b)) if a.as_char() == b.as_char() => true,
63 (SingleToken::Literal(a), TokenTree::Literal(b)) if a.to_string() == b.to_string() => {
64 true
65 }
66 _ => false,
67 }
68 }
69}
70
71pub fn find_and_replace(
72 needle: impl Into<SingleToken>,
73 replacement: TokenStream,
74 haystack: TokenStream,
75) -> TokenStream {
76 fn far(needle: &SingleToken, replacement: &TokenStream, haystack: TokenStream) -> TokenStream {
77 let mut output = TokenStream::new();
78
79 for tt in haystack {
80 let tt = match tt {
81 TokenTree::Group(g) => TokenTree::Group(Group::new(
82 g.delimiter(),
83 far(needle, replacement, g.stream()),
84 )),
85 tt if needle == &tt => {
86 output.extend(replacement.clone());
87 continue;
88 }
89 tt => tt,
90 };
91
92 output.append(tt);
93 }
94
95 output
96 }
97
98 let needle = needle.into();
99
100 far(&needle, &replacement, haystack)
101}