Skip to main content

annotate_derive/
lib.rs

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}