telety_impl/
find_and_replace.rs

1//! Based on <https://crates.io/crates/macro_find_and_replace>
2use 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}