1use proc_macro2::TokenStream;
2use quote::quote;
3use syn::spanned::Spanned;
4use syn::{Item, Result, parse};
5
6pub(crate) use attributes::*;
7
8use crate::function::AnnotatedFunction;
9use crate::module::AnnotatedModule;
10
11mod attributes;
12mod function;
13mod module;
14
15#[proc_macro_attribute]
16pub fn pragma(
17 attr: proc_macro::TokenStream,
18 item: proc_macro::TokenStream,
19) -> proc_macro::TokenStream {
20 let expanded = expand(attr, item).unwrap_or_else(|error| error.to_compile_error());
21 expanded.into()
22}
23
24#[proc_macro]
25pub fn environment(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
26 let expanded = expand_environment(input).unwrap_or_else(|error| error.to_compile_error());
27 expanded.into()
28}
29
30fn expand(attr: proc_macro::TokenStream, item: proc_macro::TokenStream) -> Result<TokenStream> {
31 let source_path = source_path_of(item.clone());
32 let attributes = Attributes::parse(attr)?;
33 let item = parse::<Item>(item)?;
34
35 let expanded = match item {
36 Item::Fn(item_fn) => AnnotatedFunction::new(item_fn, attributes, source_path).expand(),
37 Item::Mod(item_mod) => AnnotatedModule::new(item_mod, attributes, source_path).expand(),
38 _ => syn::Error::new(
39 item.span(),
40 "Pragmas are not supported for this type of construct",
41 )
42 .to_compile_error(),
43 };
44
45 Ok(expanded)
46}
47
48fn expand_environment(input: proc_macro::TokenStream) -> Result<TokenStream> {
49 let annotate_path = if input.is_empty() {
50 None
51 } else {
52 Some(parse::<syn::Path>(input)?)
53 };
54 let generated_path = environment_source_path(proc_macro::Span::call_site());
55 let generated_path = syn::LitStr::new(generated_path.as_str(), proc_macro2::Span::call_site());
56
57 let expanded = if let Some(annotate_path) = annotate_path {
58 quote! {
59 mod __annotate {
60 use #annotate_path;
61 include!(concat!(env!("OUT_DIR"), "/annotate/", #generated_path));
62 pub const fn environment() -> &'static #annotate_path::Environment {
63 &__annotate::ENVIRONMENT
64 }
65 }
66
67 #[doc(hidden)]
68 #[inline(never)]
69 pub fn __ensure_linked() {
70 __annotate::__ensure_linked();
71 }
72 }
73 } else {
74 quote! {
75 #[macro_use]
76 extern crate annotate;
77 extern crate alloc;
78
79 include!(concat!(env!("OUT_DIR"), "/annotate/", #generated_path));
80 #[doc(hidden)]
81 #[inline(never)]
82 pub fn __ensure_linked() {
83 __annotate::__ensure_linked();
84 }
85 pub const fn environment() -> &'static annotate::Environment {
86 &__annotate::ENVIRONMENT
87 }
88 }
89 };
90
91 Ok(expanded)
92}
93
94fn environment_source_path(span: proc_macro::Span) -> String {
95 let source_path = std::path::PathBuf::from(span.file());
96 let manifest_root = std::path::PathBuf::from(
97 std::path::PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap())
98 .file_name()
99 .unwrap(),
100 );
101
102 if source_path.is_absolute()
103 && let Ok(manifest_dir) = std::env::var("CARGO_MANIFEST_DIR")
104 {
105 let manifest_dir = std::path::PathBuf::from(manifest_dir);
106 if let Ok(relative_path) = source_path.strip_prefix(&manifest_dir) {
107 return manifest_root
108 .join(relative_path)
109 .to_string_lossy()
110 .replace('\\', "/");
111 }
112 }
113
114 if source_path
115 .components()
116 .next()
117 .map(|component| component.as_os_str() == manifest_root.as_os_str())
118 .unwrap_or(false)
119 {
120 return source_path.to_string_lossy().replace('\\', "/");
121 }
122
123 manifest_root
124 .join(source_path)
125 .to_string_lossy()
126 .replace('\\', "/")
127}
128
129fn source_path_of(stream: proc_macro::TokenStream) -> String {
130 let span = stream
131 .into_iter()
132 .next()
133 .map(|token| token.span())
134 .unwrap_or_else(proc_macro::Span::call_site);
135 environment_source_path(span)
136}