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