proc_debug/
lib.rs

1#![doc = include_str!("README.md")]
2
3use argp::FromArgs;
4use bat::PrettyPrinter;
5
6/// See module-level documentation
7pub use proc_debug_macro::proc_debug;
8use proc_macro2::{TokenStream, TokenTree};
9use quote::{quote, quote_spanned, ToTokens};
10use std::collections::VecDeque;
11use std::{io::Write, str::FromStr};
12use syn::*;
13use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};
14
15fn print<R>(f: impl FnOnce(&mut StandardStream) -> R) -> R {
16    let mut stdout = StandardStream::stdout(ColorChoice::Always);
17    stdout
18        .set_color(
19            ColorSpec::new()
20                .set_bg(Some(Color::Cyan))
21                .set_fg(Some(Color::Black))
22                .set_bold(true),
23        )
24        .unwrap();
25    f(&mut stdout)
26}
27
28enum MacroOutput {
29    Expr(Expr),
30    Type(Type),
31    ImplItem(Vec<ImplItem>),
32    TraitItem(Vec<TraitItem>),
33    ForeignItem(Vec<ForeignItem>),
34    Item(Vec<Item>),
35    Stmt(Vec<Stmt>),
36    Other(TokenStream),
37}
38
39#[derive(Debug, Clone, PartialEq, Eq, Hash)]
40enum MacroKind {
41    Function,
42    Attribute,
43    Derive,
44    Other,
45}
46
47impl FromStr for MacroKind {
48    type Err = &'static str;
49
50    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
51        match s {
52            "function" => Ok(Self::Function),
53            "attribute" => Ok(Self::Attribute),
54            "derive" => Ok(Self::Derive),
55            _ => Ok(Self::Other),
56        }
57    }
58}
59
60impl ToTokens for MacroOutput {
61    fn to_tokens(&self, tokens: &mut TokenStream) {
62        let rhs = match self {
63            MacroOutput::Expr(expr) => quote!(#expr),
64            MacroOutput::Type(ty) => quote!(#ty),
65            MacroOutput::ImplItem(o) => quote!(#(#o)*),
66            MacroOutput::ForeignItem(o) => quote!(#(#o)*),
67            MacroOutput::TraitItem(o) => quote!(#(#o)*),
68            MacroOutput::Item(o) => quote!(#(#o)*),
69            MacroOutput::Stmt(o) => quote!(#(#o)*),
70            MacroOutput::Other(o) => o.clone(),
71        };
72        tokens.extend(rhs);
73    }
74}
75
76fn simplify_and_replace(tokens: TokenStream, depth: usize) -> TokenStream {
77    let mut out = TokenStream::new();
78    if depth == 0 {
79        out.extend(quote!(__proc_debug_ellipsis! {}));
80        return out;
81    }
82    let mut count = 0;
83    for token in tokens {
84        match &token {
85            TokenTree::Group(g) => {
86                let inner = simplify_and_replace(g.stream(), depth - 1);
87                out.extend(Some(TokenTree::Group(proc_macro2::Group::new(
88                    g.delimiter(),
89                    inner,
90                ))));
91            }
92            TokenTree::Punct(p) if p.as_char() == ';' => {
93                count += 1;
94                out.extend(Some(token.clone()));
95                if count >= depth {
96                    out.extend(quote_spanned!(p.span() => __proc_debug_ellipsis!{}));
97                    break;
98                }
99            }
100            TokenTree::Ident(ident) if &ident.to_string() == "$crate" => {
101                out.extend(quote_spanned!(ident.span() => __proc_debug_dollar_crate!{}));
102            }
103            _ => {
104                out.extend(Some(token));
105            }
106        }
107    }
108    out
109}
110
111fn unreplace(tokens: TokenStream) -> TokenStream {
112    let mut out = TokenStream::new();
113    let mut tokens: VecDeque<_> = tokens.into_iter().collect();
114    while let Some(token) = tokens.pop_front() {
115        if let TokenTree::Ident(ident) = token.clone() {
116            match tokens.pop_front() {
117                Some(TokenTree::Punct(p)) if p.as_char() == '!' => {
118                    match tokens.pop_front() {
119                        Some(TokenTree::Group(g))
120                            if g.delimiter() == proc_macro2::Delimiter::Brace =>
121                        {
122                            match ident.to_string().as_str() {
123                                "__proc_debug_ellipsis" => {
124                                    out.extend(quote_spanned!(ident.span() => ...));
125                                    continue;
126                                }
127                                "__proc_debug_dollar_crate" => {
128                                    out.extend(quote_spanned!(ident.span() => $crate));
129                                    continue;
130                                }
131                                _ => (),
132                            }
133                            tokens.push_front(TokenTree::Group(g));
134                        }
135                        Some(o) => tokens.push_front(o),
136                        None => (),
137                    }
138                    tokens.push_front(TokenTree::Punct(p));
139                }
140                Some(o) => tokens.push_front(o),
141                None => (),
142            }
143        }
144        if let TokenTree::Group(g) = token.clone() {
145            let ng = proc_macro2::Group::new(g.delimiter(), unreplace(g.stream()));
146            out.extend(Some(TokenTree::Group(ng)));
147        } else {
148            out.extend(Some(token));
149        }
150    }
151    out
152}
153
154impl MacroOutput {
155    fn from_tokens(tokens: TokenStream, kind: &MacroKind) -> Self {
156        struct Sequentary<T>(Vec<T>);
157        impl<T> syn::parse::Parse for Sequentary<T>
158        where
159            T: syn::parse::Parse,
160        {
161            fn parse(input: parse::ParseStream) -> Result<Self> {
162                let mut v = Vec::new();
163                while !input.is_empty() {
164                    v.push(input.parse()?)
165                }
166                Ok(Self(v))
167            }
168        }
169        if kind == &MacroKind::Function {
170            if let Ok(ident) = parse2::<Ident>(tokens.clone()) {
171                if ident.to_string().chars().next().unwrap().is_uppercase() {
172                    return Self::Type(parse_quote! {#ident});
173                } else {
174                    return Self::Expr(parse_quote!(#ident));
175                }
176            }
177            if let Ok(ty) = parse2::<Type>(tokens.clone()) {
178                return Self::Type(ty);
179            }
180        }
181        if let Ok(s) = parse2::<Sequentary<_>>(tokens.clone()) {
182            return Self::ImplItem(s.0);
183        }
184        if let Ok(s) = parse2::<Sequentary<_>>(tokens.clone()) {
185            return Self::TraitItem(s.0);
186        }
187        if let Ok(s) = parse2::<Sequentary<_>>(tokens.clone()) {
188            return Self::ForeignItem(s.0);
189        }
190        if let Ok(s) = parse2::<Sequentary<_>>(tokens.clone()) {
191            return Self::Item(s.0);
192        }
193        if let Ok(s) = parse2::<Sequentary<_>>(tokens.clone()) {
194            return Self::Stmt(s.0);
195        }
196        Self::Other(tokens)
197    }
198
199    fn emit(&self) -> TokenStream {
200        match self {
201            MacroOutput::Expr(expr) => quote! {#expr},
202            MacroOutput::Type(ty) => quote! {#ty},
203            o => quote! {#o},
204        }
205    }
206}
207
208fn show_macro_call(
209    modpath: &str,
210    macro_name: &str,
211    file: &str,
212    line: usize,
213    macro_kind: &str,
214    macro_inputs: &[String],
215) {
216    let content = match macro_kind {
217        "function" => format!("{macro_name}!{{{}}}", macro_inputs[0]),
218        "attribute" => format!(
219            "#[{}({})]\n{}",
220            macro_name, macro_inputs[0], macro_inputs[1]
221        ),
222        "derive" => format!("#[derive({})]\n{}", macro_inputs[0], macro_inputs[1]),
223        _ => format!("{}", macro_inputs.join(",")),
224    };
225    let content = content
226        .split("\n")
227        .map(|s| format!("  {}", s))
228        .collect::<Vec<_>>()
229        .join("\n");
230    print(|out| writeln!(out, "👉 input of {modpath}::{macro_name} ({file}:{line})",)).unwrap();
231    PrettyPrinter::new()
232        .input_from_reader(content.as_bytes())
233        .language("rust")
234        .print()
235        .unwrap();
236    writeln!(std::io::stdout(), "",).unwrap();
237}
238
239fn show_macro_output(modpath: &str, macro_name: &str, file: &str, line: usize, macro_output: &str) {
240    print(|out| writeln!(out, "👉 output of {modpath}::{macro_name} ({file}:{line})",)).unwrap();
241    let content = macro_output
242        .split("\n")
243        .map(|s| format!("  {}", s))
244        .collect::<Vec<_>>()
245        .join("\n");
246    PrettyPrinter::new()
247        .input_from_bytes(content.as_bytes())
248        .language("rust")
249        .print()
250        .unwrap();
251    writeln!(std::io::stdout(), "",).unwrap();
252}
253
254/// Input for `proc-debug`
255#[derive(FromArgs)]
256struct ProcDebugArgs {
257    /// debug all macros
258    #[argp(switch, short = 'a')]
259    all: bool,
260    /// hide outputs match
261    #[argp(option, short = 'n')]
262    not: Vec<String>,
263    /// full or partial path of macro definition
264    #[argp(option, short = 'p')]
265    path: Vec<String>,
266    /// search queries to show debug
267    #[argp(positional, greedy)]
268    queries: Vec<String>,
269    /// depth to show in macro output
270    #[argp(option, short = 'd')]
271    depth: Option<usize>,
272    /// verbose
273    #[argp(switch, short = 'v')]
274    verbose: bool,
275}
276
277#[test]
278fn test_split_args() {
279    assert_eq!(
280        split_args(r#"  --all  -a  ' b c ' " -d '' "  "#),
281        vec![
282            "--all".to_owned(),
283            "-a".to_owned(),
284            " b c ".to_owned(),
285            " -d '' ".to_owned()
286        ]
287    );
288}
289
290fn split_args(s: &str) -> Vec<String> {
291    let mut it = s.chars().fuse();
292    let mut res = Vec::new();
293    let mut r = String::new();
294    while let Some(c) = it.next() {
295        match c {
296            '\\' => {
297                if let Some(c) = it.next() {
298                    r.push(c);
299                }
300            }
301            '"' | '\'' => {
302                let delim = c;
303                while let Some(c) = it.next() {
304                    match c {
305                        '\\' => {
306                            if let Some(c) = it.next() {
307                                r.push(c);
308                            }
309                        }
310                        c if c == delim => {
311                            res.push(r);
312                            r = String::new();
313                            break;
314                        }
315                        c => r.push(c),
316                    }
317                }
318            }
319            c if c.is_ascii_whitespace() => {
320                if r.len() > 0 {
321                    res.push(r);
322                    r = String::new();
323                }
324            }
325            c => r.push(c),
326        }
327    }
328    if r.len() > 0 {
329        res.push(r);
330    }
331    res
332}
333
334impl ProcDebugArgs {
335    fn from_env() -> Option<Self> {
336        let flags = std::env::var("PROC_DEBUG_FLAGS").ok()?;
337        let flags = split_args(&flags);
338        Some(
339            ProcDebugArgs::from_args(&["proc-debug"], &flags).unwrap_or_else(|early_exit| {
340                let mut stderr = StandardStream::stderr(ColorChoice::Always);
341                stderr
342                    .set_color(
343                        ColorSpec::new()
344                            .set_bg(Some(Color::Yellow))
345                            .set_fg(Some(Color::Black))
346                            .set_bold(true),
347                    )
348                    .unwrap();
349                match early_exit {
350                    argp::EarlyExit::Help(help) => {
351                        writeln!(&mut stderr, "{}", help.generate_default()).unwrap()
352                    }
353                    argp::EarlyExit::Err(err) => writeln!(
354                        &mut stderr,
355                        "{} \n\n Set PROC_DEBUG_FLAGS=\"--help\" for more information.",
356                        err
357                    )
358                    .unwrap(),
359                }
360                std::process::exit(1)
361            }),
362        )
363    }
364}
365
366#[allow(unused)]
367struct Entry<'a> {
368    label: &'a str,
369    file: &'a str,
370    line: usize,
371    modpath: &'a str,
372    macro_kind: &'a str,
373    macro_name: &'a str,
374    macro_inputs: &'a [String],
375}
376
377impl<'a> Entry<'a> {
378    fn check_filter(&self, args: &ProcDebugArgs) -> bool {
379        let content = [&self.label, &self.file, &self.modpath, &self.macro_name];
380        let pattern = format!("{}::{}", &self.modpath, &self.macro_name);
381
382        if args.all {
383            return true;
384        }
385        if content
386            .iter()
387            .any(|s| args.not.iter().any(|t| s.contains(t)))
388        {
389            return false;
390        }
391        if args.path.iter().any(|m| {
392            m == &pattern
393                || pattern.starts_with(&format!("{}::", m))
394                || pattern.ends_with(&format!("::{}", m))
395        }) {
396            return true;
397        }
398        if content
399            .iter()
400            .any(|s| args.queries.iter().any(|t| s.contains(t)))
401        {
402            return true;
403        }
404        false
405    }
406}
407
408#[doc(hidden)]
409pub fn proc_wrapper<F: FnOnce() -> TokenStream>(
410    label: &str,
411    file: &str,
412    line: usize,
413    modpath: &str,
414    macro_kind: &str,
415    macro_name: &str,
416    macro_inputs: &[String],
417    f: F,
418) -> TokenStream {
419    let entry = Entry {
420        label,
421        file,
422        line,
423        modpath,
424        macro_kind,
425        macro_name,
426        macro_inputs,
427    };
428    let ret = f();
429    if let Some(args) = ProcDebugArgs::from_env() {
430        if entry.check_filter(&args) {
431            show_macro_call(modpath, macro_name, file, line, macro_kind, macro_inputs);
432            let tokens: TokenStream = ret.into();
433            let output =
434                MacroOutput::from_tokens(tokens.clone(), &MacroKind::from_str(macro_kind).unwrap());
435            let simplified = simplify_and_replace(
436                tokens,
437                if args.verbose {
438                    usize::MAX
439                } else {
440                    args.depth.unwrap_or(4)
441                },
442            );
443
444            show_macro_output(
445                modpath,
446                macro_name,
447                file,
448                line,
449                &unreplace(simplified).to_string(),
450            );
451            output.emit().into()
452        } else {
453            ret
454        }
455    } else {
456        ret
457    }
458}
proc_debug/lib.rs

proc_debug/
lib.rs
