1use proc_macro2::TokenStream;
6use proc_macro_error2::abort;
7use quote::{quote, ToTokens};
8use syn::{parse_quote, spanned::Spanned, ItemFn};
9
10pub fn carbond_track_macro(input: TokenStream) -> TokenStream {
11 let span = input.span();
12 let mut input: ItemFn = match syn::parse2(input) {
13 Ok(input) => input,
14 Err(_) => {
15 abort!(span, "#[carbond_track] must be used on a function.")
16 }
17 };
18
19 let ty: Box<syn::Type> = extract_type(&input.sig);
21 input.sig.output = parse_quote! { -> carbond_lib::data::carbon_tracking::CarbonTracking<#ty> };
22
23 let block = input.block.to_token_stream();
24
25 let signature = input.sig;
26
27 quote! {
28 #signature {
29 use carbond_lib::rand::seq::SliceRandom;
30 use rand::prelude::IndexedRandom;
31 let core = carbond_lib::core_affinity::get_core_ids().unwrap().choose(&mut rand::rng()).unwrap().id;
32 carbond_lib::core_affinity::set_for_current(carbond_lib::core_affinity::CoreId { id: core });
33 let mut core = core as u32;
34 let cpu_cycles_pre = unsafe {
35 core::arch::x86_64::__rdtscp(&mut core)
36 };
37 let val = || #block;
38 let val = val();
39 let cpu_track = {
40 let cycles = unsafe {
41 core::arch::x86_64::__rdtscp(&mut core)
42 } - cpu_cycles_pre;
43 carbond_lib::data::carbon_tracking::CpuTracking::new(core, cycles)
44 };
45 return carbond_lib::data::carbon_tracking::CarbonTracking::new(val, cpu_track);
46 }
47 }
48}
49
50fn extract_type(signature: &syn::Signature) -> Box<syn::Type> {
51 if let syn::ReturnType::Type(_, ref ty) = signature.output {
52 ty.clone()
53 } else {
54 parse_quote!(())
55 }
56}
57
58#[cfg(test)]
59mod tests {
60 use crate::macros::carbond_track_macro;
61 use assert_tokenstreams_eq::assert_tokenstreams_eq;
62 use quote::{quote, ToTokens};
63
64 #[test]
65 fn test_carbond_track_macro_explicit_return() {
66 let input = quote! {
67 fn foo() -> i64 {
68 let a = 6424_i64;
69 let b = a + 500;
70 return b;
71 }
72 };
73 let result = carbond_track_macro(input);
74 let expected = quote! {
75 fn foo() -> carbond_lib::data::carbon_tracking::CarbonTracking<i64> {
76 use carbond_lib::rand::seq::SliceRandom;
77 use rand::prelude::IndexedRandom;
78 let core = carbond_lib::core_affinity::get_core_ids().unwrap().choose(&mut rand::rng()).unwrap().id;
79 carbond_lib::core_affinity::set_for_current(carbond_lib::core_affinity::CoreId { id: core });
80 let mut core = core as u32;
81 let cpu_cycles_pre = unsafe { core::arch::x86_64::__rdtscp(&mut core) };
82 let val = || {
83 let a = 6424_i64;
84 let b = a + 500;
85 return b;
86 };
87 let val = val();
88 let cpu_track = {
89 let cycles = unsafe { core::arch::x86_64::__rdtscp(&mut core) } - cpu_cycles_pre;
90 carbond_lib::data::carbon_tracking::CpuTracking::new(core, cycles)
91 };
92 return carbond_lib::data::carbon_tracking::CarbonTracking::new(val, cpu_track);
93 }
94 };
95 assert_tokenstreams_eq!(&expected, &result.to_token_stream());
96 }
97
98 #[test]
99 fn test_carbond_track_macro_implicit_return() {
100 let input = quote! {
101 fn foo() -> i64 {
102 let a = 6424_i64;
103 a + 500
104 }
105 };
106 let result = carbond_track_macro(input);
107 let expected = quote! {
108 fn foo() -> carbond_lib::data::carbon_tracking::CarbonTracking<i64> {
109 use carbond_lib::rand::seq::SliceRandom;
110 use rand::prelude::IndexedRandom;
111 let core = carbond_lib::core_affinity::get_core_ids().unwrap().choose(&mut rand::rng()).unwrap().id;
112 carbond_lib::core_affinity::set_for_current(carbond_lib::core_affinity::CoreId { id: core });
113 let mut core = core as u32;
114 let cpu_cycles_pre = unsafe { core::arch::x86_64::__rdtscp(&mut core) };
115 let val = || {
116 let a = 6424_i64;
117 a + 500
118 };
119 let val = val();
120 let cpu_track = {
121 let cycles = unsafe { core::arch::x86_64::__rdtscp(&mut core) } - cpu_cycles_pre;
122 carbond_lib::data::carbon_tracking::CpuTracking::new(core, cycles)
123 };
124 return carbond_lib::data::carbon_tracking::CarbonTracking::new(val, cpu_track);
125 }
126 };
127 assert_tokenstreams_eq!(&expected, &result.to_token_stream());
128 }
129
130 #[test]
131 fn test_carbond_track_macro_unit_return() {
132 let input = quote! {
133 fn foo() -> i64 {
134 let a = 6424_i64;
135 let b = a + 500;
136 }
137 };
138 let result = carbond_track_macro(input);
139 let expected = quote! {
140 fn foo() -> carbond_lib::data::carbon_tracking::CarbonTracking<i64> {
141 use carbond_lib::rand::seq::SliceRandom;
142 use rand::prelude::IndexedRandom;
143 let core = carbond_lib::core_affinity::get_core_ids().unwrap().choose(&mut rand::rng()).unwrap().id;
144 carbond_lib::core_affinity::set_for_current(carbond_lib::core_affinity::CoreId { id: core });
145 let mut core = core as u32;
146 let cpu_cycles_pre = unsafe { core::arch::x86_64::__rdtscp(&mut core) };
147 let val = || {
148 let a = 6424_i64;
149 let b = a + 500;
150 };
151 let val = val();
152 let cpu_track = {
153 let cycles = unsafe { core::arch::x86_64::__rdtscp(&mut core) } - cpu_cycles_pre;
154 carbond_lib::data::carbon_tracking::CpuTracking::new(core, cycles)
155 };
156 return carbond_lib::data::carbon_tracking::CarbonTracking::new(val, cpu_track);
157 }
158 };
159 assert_tokenstreams_eq!(&expected, &result.to_token_stream());
160 }
161}