flux_attrs_impl/
lib.rs

1mod ast;
2mod extern_spec;
3
4use proc_macro2::{Ident, TokenStream, TokenTree};
5use quote::{ToTokens, quote, quote_spanned};
6use syn::{
7    Attribute, ItemEnum, ItemStruct, Token, bracketed, parse::ParseStream, parse_quote,
8    spanned::Spanned,
9};
10
11pub const FLUX_ATTRS: &[&str] = &[
12    "assoc",
13    "field",
14    "generics",
15    "invariant",
16    "opaque",
17    "reflect",
18    "refined_by",
19    "sig",
20    "trusted",
21    "trusted_impl",
22    "variant",
23    "should_fail",
24    "opts",
25    "reft",
26];
27
28pub fn extern_spec(attr: TokenStream, tokens: TokenStream) -> TokenStream {
29    extern_spec::transform_extern_spec(attr, tokens).unwrap_or_else(|err| err.to_compile_error())
30}
31
32pub fn flux_tool_item_attr(name: &str, attr: TokenStream, item: TokenStream) -> TokenStream {
33    // I don't really know what I'm doing here, but spanning the quote with the item's span seems
34    // to behave correctly.
35    let span = item.span();
36    let name = TokenTree::Ident(Ident::new(name, span));
37    if attr.is_empty() {
38        quote_spanned! {span=>
39            #[flux_tool::#name]
40            #item
41        }
42    } else {
43        quote_spanned! {span=>
44            #[flux_tool::#name(#attr)]
45            #item
46        }
47    }
48}
49
50pub fn refined_by(attr: TokenStream, item: TokenStream) -> TokenStream {
51    let span = item.span();
52    let mut item = match syn::parse2::<syn::Item>(item) {
53        Ok(item) => item,
54        Err(err) => return err.to_compile_error(),
55    };
56
57    match &mut item {
58        syn::Item::Enum(item_enum) => refined_by_enum(item_enum),
59        syn::Item::Struct(item_struct) => refined_by_struct(item_struct),
60        _ => return syn::Error::new(span, "expected struct or enum").to_compile_error(),
61    }
62
63    if cfg!(flux_sysroot) {
64        quote_spanned! {span=>
65            #[flux_tool::refined_by(#attr)]
66            #item
67        }
68    } else {
69        item.to_token_stream()
70    }
71}
72
73fn refined_by_enum(item_enum: &mut ItemEnum) {
74    for variant in &mut item_enum.variants {
75        flux_tool_attrs(&mut variant.attrs);
76    }
77}
78
79fn refined_by_struct(item_struct: &mut ItemStruct) {
80    for field in &mut item_struct.fields {
81        flux_tool_attrs(&mut field.attrs);
82    }
83}
84
85fn flux_tool_attrs(attrs: &mut Vec<Attribute>) {
86    if cfg!(flux_sysroot) {
87        for attr in attrs {
88            transform_flux_attr(attr);
89        }
90    } else {
91        attrs.retain(|attr| !is_flux_attr(attr));
92    }
93}
94
95fn path_is_one_of(path: &syn::Path, idents: &[&str]) -> bool {
96    idents.iter().any(|ident| path.is_ident(ident))
97}
98
99fn is_flux_attr(attr: &syn::Attribute) -> bool {
100    let path = attr.path();
101    if path.segments.len() >= 2 {
102        let ident = &path.segments[0].ident;
103        ident == "flux" || ident == "flux_rs"
104    } else {
105        path_is_one_of(path, FLUX_ATTRS)
106    }
107}
108
109fn transform_flux_attr(attr: &mut syn::Attribute) {
110    let path = path_of_attr_mut(attr);
111    if path.leading_colon.is_some() {
112        return;
113    }
114    if path.segments.len() >= 2 {
115        let ident = &mut path.segments[0].ident;
116        if ident == "flux" || ident == "flux_rs" {
117            *ident = Ident::new("flux_tool", ident.span());
118        }
119        return;
120    } else if path_is_one_of(path, FLUX_ATTRS) {
121        *path = parse_quote!(flux_tool::#path);
122    }
123}
124
125fn path_of_attr_mut(attr: &mut Attribute) -> &mut syn::Path {
126    match &mut attr.meta {
127        syn::Meta::Path(path) => path,
128        syn::Meta::List(metalist) => &mut metalist.path,
129        syn::Meta::NameValue(namevalue) => &mut namevalue.path,
130    }
131}
132
133pub fn flux(tokens: TokenStream) -> TokenStream {
134    syn::parse2::<ast::Items>(tokens)
135        .map_or_else(|err| err.to_compile_error(), ToTokens::into_token_stream)
136}
137
138pub fn defs(tokens: TokenStream) -> TokenStream {
139    quote! {
140        #[flux::defs { #tokens }]
141        const _: () = {};
142    }
143}
144
145pub fn tokens_or_default<T: ToTokens + Default>(x: Option<&T>, tokens: &mut TokenStream) {
146    match x {
147        Some(t) => t.to_tokens(tokens),
148        None => T::default().to_tokens(tokens),
149    }
150}
151
152fn parse_inner(input: ParseStream, attrs: &mut Vec<Attribute>) -> syn::Result<()> {
153    while input.peek(Token![#]) && input.peek2(Token![!]) {
154        attrs.push(input.call(single_parse_inner)?);
155    }
156    Ok(())
157}
158
159fn single_parse_inner(input: ParseStream) -> syn::Result<Attribute> {
160    let content;
161    Ok(Attribute {
162        pound_token: input.parse()?,
163        style: syn::AttrStyle::Inner(input.parse()?),
164        bracket_token: bracketed!(content in input),
165        meta: content.parse()?,
166    })
167}
168
169fn outer(attrs: &[Attribute]) -> impl Iterator<Item = &Attribute> {
170    fn is_outer(attr: &&Attribute) -> bool {
171        match attr.style {
172            syn::AttrStyle::Outer => true,
173            syn::AttrStyle::Inner(_) => false,
174        }
175    }
176    attrs.iter().filter(is_outer)
177}
178
179fn inner(attrs: &[Attribute]) -> impl Iterator<Item = &Attribute> {
180    fn is_inner(attr: &&Attribute) -> bool {
181        match attr.style {
182            syn::AttrStyle::Outer => false,
183            syn::AttrStyle::Inner(_) => true,
184        }
185    }
186    attrs.iter().filter(is_inner)
187}