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    } 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}