thag_proc_macros 1.0.1

Procedural macros for the thag_rs and thag_profiler crates
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109

#![allow(clippy::module_name_repetitions)]
use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{
    parse::{Parse, ParseStream},
    parse_macro_input,
    punctuated::Punctuated,
    Ident, Result, Token,
};

/// Arguments for the `profile` macro
struct ProfileArgs {
    name: String,
    args: Punctuated<Ident, Token![,]>,
}

impl Parse for ProfileArgs {
    fn parse(input: ParseStream) -> Result<Self> {
        // Parse the profile name as an identifier
        let name = input.parse::<Ident>()?.to_string();

        // Parse remaining arguments as identifiers separated by commas
        let args = if input.is_empty() {
            Punctuated::new()
        } else {
            input.parse::<Token![,]>()?;
            input.parse_terminated(Ident::parse, Token![,])?
        };

        Ok(Self { name, args })
    }
}

pub fn profile_impl(input: TokenStream) -> TokenStream {
    let ProfileArgs { name, args } = parse_macro_input!(input as ProfileArgs);

    // Extract flags from args
    let has_time = args.iter().any(|arg| arg == "time" || arg == "both");
    let has_mem_summary = args.iter().any(|arg| arg == "mem_summary" || arg == "both");
    let has_mem_detail = args.iter().any(|arg| arg == "mem_detail");
    let is_async = args.iter().any(|arg| arg == "async_fn");
    let is_unbounded = args.iter().any(|arg| arg == "unbounded");
    let use_global = args.iter().any(|arg| arg == "global");

    // Determine profile type
    let profile_type = if use_global {
        quote! { ::thag_profiler::get_global_profile_type() }
    } else if has_time && (has_mem_summary || has_mem_detail) {
        quote! { ::thag_profiler::ProfileType::Both }
    } else if has_time {
        quote! { ::thag_profiler::ProfileType::Time }
    } else if has_mem_summary || has_mem_detail {
        quote! { ::thag_profiler::ProfileType::Memory }
    } else {
        quote! { ::thag_profiler::get_global_profile_type() }
    };

    // Determine if detailed memory is enabled
    let detailed_memory = has_mem_detail;

    // Determine line numbers
    let (start_line, end_line) = if has_time && !(has_mem_summary || has_mem_detail) {
        // Time only - no line numbers needed
        (quote! { None }, quote! { None })
    } else if is_unbounded {
        // Memory with unbounded - only start line
        (quote! { Some(line!()) }, quote! { None })
    } else {
        // Memory with bounded - need end marker
        // Convert the string to an identifier
        let end_line_fn = format_ident!("end_{name}");
        (quote! { Some(line!()) }, quote! { Some(#end_line_fn()) })
    };

    let profile_id = format_ident!("{name}");

    // Generate the profile creation code
    #[cfg(not(feature = "full_profiling"))]
    let expanded = quote! {
        let #profile_id = ::thag_profiler::Profile::new(
            Some(#name),
            None,
            #profile_type,
            #is_async,
            #detailed_memory,
            file!(),
            #start_line,
            #end_line
        );
    };

    #[cfg(feature = "full_profiling")]
    let expanded = quote! {
        let #profile_id = ::thag_profiler::safe_alloc! {
            ::thag_profiler::Profile::new(
                Some(#name),
                None,
                #profile_type,
                #is_async,
                #detailed_memory,
                file!(),
                #start_line,
                #end_line
            )
        };
    };

    expanded.into()
}