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 } else {
68 quote! {
69 #[macro_use]
70 extern crate annotate;
71 extern crate alloc;
72
73 include!(concat!(env!("OUT_DIR"), "/annotate/", #generated_path));
74 pub const fn environment() -> &'static annotate::Environment {
75 &__annotate::ENVIRONMENT
76 }
77 }
78 };
79
80 Ok(expanded)
81}
82
83fn environment_source_path(span: proc_macro::Span) -> String {
84 let source_path = std::path::PathBuf::from(span.file());
85 let manifest_root = std::path::PathBuf::from(
86 std::path::PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap())
87 .file_name()
88 .unwrap(),
89 );
90
91 if source_path.is_absolute()
92 && let Ok(manifest_dir) = std::env::var("CARGO_MANIFEST_DIR")
93 {
94 let manifest_dir = std::path::PathBuf::from(manifest_dir);
95 if let Ok(relative_path) = source_path.strip_prefix(&manifest_dir) {
96 return manifest_root
97 .join(relative_path)
98 .to_string_lossy()
99 .replace('\\', "/");
100 }
101 }
102
103 if source_path
104 .components()
105 .next()
106 .map(|component| component.as_os_str() == manifest_root.as_os_str())
107 .unwrap_or(false)
108 {
109 return source_path.to_string_lossy().replace('\\', "/");
110 }
111
112 manifest_root.join(source_path).to_string_lossy().replace('\\', "/")
113}
114
115fn source_path_of(stream: proc_macro::TokenStream) -> String {
116 let span = stream
117 .into_iter()
118 .next()
119 .map(|token| token.span())
120 .unwrap_or_else(proc_macro::Span::call_site);
121 environment_source_path(span)
122}