quirky_binder_codegen_macro 0.1.0

Quirky Binder Codegen Macro
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

use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, ItemFn};

const MONITOR_OPTION: &str = "quirky_binder_monitor";

pub fn tracking_allocator_static(_input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    quote! {
        #[cfg(feature = #MONITOR_OPTION)]
        #[global_allocator]
        static GLOBAL: tracking_allocator::Allocator<std::alloc::System> = tracking_allocator::Allocator::system();

        #[cfg(feature = #MONITOR_OPTION)]
        mod quirky_binder_monitor {
            use std::sync::Mutex;

            use tracking_allocator::{AllocationGroupId, AllocationTracker};

            pub struct MemoryTracker;

            pub struct MemoryTrackerData {
                pub object: usize,
                pub wrapped: usize,
            }

            pub static DATA: Mutex<MemoryTrackerData> = Mutex::new(MemoryTrackerData {
                object: 0,
                wrapped: 0,
            });

            impl AllocationTracker for MemoryTracker {
                fn allocated(
                    &self,
                    _addr: usize,
                    object_size: usize,
                    wrapped_size: usize,
                    _group_id: AllocationGroupId,
                    ) {
                    let mut data = DATA.lock().unwrap();
                    data.object += object_size;
                    data.wrapped += wrapped_size;
                }

                fn deallocated(
                    &self,
                    _addr: usize,
                    object_size: usize,
                    wrapped_size: usize,
                    _source_group_id: AllocationGroupId,
                    _current_group_id: AllocationGroupId,
                    ) {
                    let mut data = DATA.lock().unwrap();
                    data.object -= object_size;
                    data.wrapped -= wrapped_size;
                }
            }
        }
    }
    .into()
}

pub fn tracking_allocator_main(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let ItemFn {
        attrs,
        vis,
        sig,
        block,
    } = parse_macro_input!(item as ItemFn);

    quote! {
        #[cfg(not(feature = #MONITOR_OPTION))]
        #(#attrs)*
        #vis #sig {
            let ret = #block;

            ret
        }

        #[cfg(feature = #MONITOR_OPTION)]
        #(#attrs)*
        #vis #sig {
            use tracking_allocator::{AllocationGroupToken, AllocationRegistry, Allocator};

            AllocationRegistry::set_global_tracker(quirky_binder_monitor::MemoryTracker)
                .expect("no other global tracker should be set yet");

            let mut local_token =
                AllocationGroupToken::register().expect("failed to register allocation group");

            let local_guard = local_token.enter();

            AllocationRegistry::enable_tracking();

            let ret = #block;

            AllocationRegistry::disable_tracking();

            ret
        }
    }
    .into()
}