Skip to main content

flux_attrs_impl/
lib.rs

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