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