hax_rust_engine_macros/
lib.rs

1//! Helper crate providing procedural macros for the Rust engine of hax.
2//!
3//! Currently it provides the following.
4//!  - Macros for deriving groups of traits.
5//!    Most of the type from the AST have the same bounds, so that helps deduplicating a lot.
6//!    Also, the fact those derive groups are named is helpful: for instance for code generation
7//!    a simple `use derive_group_for_ast_base as derive_group_for_ast` can change what is to be
8//!    derived without any attribute manipulation.
9
10use proc_macro::TokenStream;
11use proc_macro2::{Group, Ident, Span};
12use quote::{ToTokens, quote};
13use syn::{
14    Field, FieldsUnnamed, Token, parse_macro_input, parse_quote, punctuated::Punctuated,
15    token::Paren, visit_mut::VisitMut,
16};
17use utils::*;
18
19mod partial_application;
20mod replace;
21
22mod utils {
23    use super::*;
24    pub(crate) fn crate_name() -> Ident {
25        let krate = module_path!().split("::").next().unwrap();
26        Ident::new(krate, Span::call_site())
27    }
28
29    /// Prepends a `proc_macro2::TokenStream` to a `TokenStream`
30    pub(crate) fn prepend(item: TokenStream, prefix: proc_macro2::TokenStream) -> TokenStream {
31        let item: proc_macro2::TokenStream = item.into();
32        quote! {
33            #prefix
34            #item
35        }
36        .into()
37    }
38
39    /// Add a derive attribute to `item`
40    pub(crate) fn add_derive(item: TokenStream, payload: proc_macro2::TokenStream) -> TokenStream {
41        prepend(item, quote! {#[derive(#payload)]})
42    }
43}
44
45/// Derive the common derives for the hax engine AST.
46/// This is a equivalent to `derive_group_for_ast_serialization` and `derive_group_for_ast_base`.
47#[proc_macro_attribute]
48pub fn derive_group_for_ast(_attr: TokenStream, item: TokenStream) -> TokenStream {
49    let krate = crate_name();
50    prepend(
51        item,
52        quote! {
53            #[#krate::derive_group_for_ast_base]
54            #[#krate::derive_group_for_ast_serialization]
55        },
56    )
57}
58
59/// Derive the necessary [de]serialization related traits for nodes in the AST.
60#[proc_macro_attribute]
61pub fn derive_group_for_ast_serialization(_attr: TokenStream, item: TokenStream) -> TokenStream {
62    add_derive(
63        item,
64        quote! {::serde::Deserialize, ::serde::Serialize, ::schemars::JsonSchema},
65    )
66}
67
68/// Derive the basic necessary traits for nodes in the AST.
69#[proc_macro_attribute]
70pub fn derive_group_for_ast_base(_attr: TokenStream, item: TokenStream) -> TokenStream {
71    add_derive(
72        item,
73        quote! {Debug, Clone, Hash, Eq, PartialEq, PartialOrd, Ord, derive_generic_visitor::Drive, derive_generic_visitor::DriveMut},
74    )
75}
76
77/// Adds a new field with a fresh name to an existing `struct` type definition.
78/// The new field contains error handling and span information to be used with a
79/// visitor. This macro will also derive implementations of
80/// [`hax_rust_engine::ast::visitors::wrappers::VisitorWithErrors`] and
81/// [`hax_rust_engine::ast::HasSpan`] for the struct.
82#[proc_macro_attribute]
83pub fn setup_error_handling_struct(_attr: TokenStream, item: TokenStream) -> TokenStream {
84    let mut item: syn::ItemStruct = parse_macro_input!(item);
85    // Deal with the case of unit structs.
86    if let fields @ syn::Fields::Unit = &mut item.fields {
87        let span = Group::new(proc_macro2::Delimiter::Brace, fields.to_token_stream()).delim_span();
88        *fields = syn::Fields::Unnamed(FieldsUnnamed {
89            paren_token: Paren { span },
90            unnamed: Punctuated::default(),
91        })
92    }
93    /// Computes a fresh identifier given a list of existing identifiers.
94    fn fresh_ident(base: &str, existing: &[Ident]) -> Ident {
95        let existing: std::collections::HashSet<_> =
96            existing.iter().map(|id| id.to_string()).collect();
97
98        (0..)
99            .map(|i| {
100                if i == 0 {
101                    base.to_string()
102                } else {
103                    format!("{}{}", base, i)
104                }
105            })
106            .find(|name| !existing.contains(name))
107            .map(|name| Ident::new(&name, Span::call_site()))
108            .expect("should always find a fresh identifier")
109    }
110    // Collect fields, disregarding their kind (are they named or not)
111    let (fields, named) = match &mut item.fields {
112        syn::Fields::Named(fields_named) => (&mut fields_named.named, true),
113        syn::Fields::Unnamed(fields_unnamed) => (&mut fields_unnamed.unnamed, false),
114        syn::Fields::Unit => unreachable!("Unit structs were dealt with."),
115    };
116
117    let existing_names = fields
118        .iter()
119        .flat_map(|f| &f.ident)
120        .cloned()
121        .collect::<Vec<_>>();
122
123    let (extra_field_ident, extra_field_ident_ts) = if named {
124        let ident = fresh_ident("error_handling_state", &existing_names);
125        (Some(ident.clone()), ident.to_token_stream())
126    } else {
127        (
128            None,
129            syn::LitInt::new(&format!("{}", fields.len()), Span::call_site()).to_token_stream(),
130        )
131    };
132
133    let krate = {
134        use proc_macro_crate::{FoundCrate, crate_name};
135        match crate_name("hax-rust-engine").unwrap() {
136            FoundCrate::Itself => quote!(crate),
137            FoundCrate::Name(name) => {
138                let ident = Ident::new(&name, Span::call_site());
139                quote!( #ident )
140            }
141        }
142    };
143
144    fields.push(Field {
145        attrs: vec![],
146        vis: syn::Visibility::Inherited,
147        mutability: syn::FieldMutability::None,
148        ident: extra_field_ident,
149        colon_token: named.then_some(Token![:](Span::call_site())),
150        ty: parse_quote! {#krate::ast::visitors::wrappers::ErrorHandlingState},
151    });
152
153    let struct_name = &item.ident;
154    let generics = &item.generics;
155    quote! {
156        #item
157        impl #generics #krate::ast::HasSpan for #struct_name #generics {
158            fn span(&self) -> #krate::ast::span::Span {
159                self.#extra_field_ident_ts.0.clone()
160            }
161            fn span_mut(&mut self) -> &mut #krate::ast::span::Span {
162                &mut self.#extra_field_ident_ts.0
163            }
164        }
165        impl #generics #krate::ast::visitors::wrappers::VisitorWithErrors for #struct_name #generics {
166            fn error_vault(&mut self) -> &mut #krate::ast::visitors::wrappers::ErrorVault {
167                &mut self.#extra_field_ident_ts.1
168            }
169        }
170    }
171    .into()
172}
173
174#[proc_macro_attribute]
175/// Replaces all occurrences of an identifier within the attached item.
176///
177/// For example, `#[replace(Name => A, B, C)]` will replace `Name` by `A, B, C`
178/// in the item the proc-macro is applied on.
179///
180/// The special case `#[replace(Name => include(VisitableAstNodes))]` will
181/// expand to a list of visitable AST nodes. This is useful in practice, as this
182/// list is often repeated.
183pub fn replace(attr: TokenStream, item: TokenStream) -> TokenStream {
184    replace::replace(attr, item)
185}
186
187/// An attribute procedural macro that creates a new `macro_rules!` definition
188/// by partially applying an existing macro or function with a given token stream.
189///
190/// Usage:
191/// ```rust,ignore
192/// #[partial_apply(original_macro!, my_expression,)]
193/// macro_rules! new_proxy_macro {
194///     // This content is ignored and replaced by the proc macro.
195/// }
196/// ```
197#[proc_macro_attribute]
198pub fn partial_apply(attr: TokenStream, item: TokenStream) -> TokenStream {
199    partial_application::partial_apply(attr, item)
200}
201
202/// Prepend the body any associated function with the given attribute payload.
203/// ```rust,ignore
204/// #[prepend_associated_functions_with(println!("self is {self}");)]
205/// impl Foo {
206///   fn f(self) {}
207/// }
208/// ```
209///
210/// Expands to:
211/// ```rust,ignore
212/// impl Foo {
213///   fn f(self) {
214///     println!("self is {self}");
215///   }
216/// }
217/// ```
218#[proc_macro_attribute]
219pub fn prepend_associated_functions_with(attr: TokenStream, item: TokenStream) -> TokenStream {
220    struct Visitor {
221        prefix: syn::Expr,
222    }
223    impl VisitMut for Visitor {
224        fn visit_item_impl_mut(&mut self, impl_block: &mut syn::ItemImpl) {
225            for item in &mut impl_block.items {
226                let syn::ImplItem::Fn(impl_item_fn) = item else {
227                    continue;
228                };
229                impl_item_fn.block.stmts.insert(
230                    0,
231                    syn::Stmt::Expr(self.prefix.clone(), Some(Token![;](Span::mixed_site()))),
232                );
233            }
234        }
235    }
236    let mut item: syn::Item = parse_macro_input!(item);
237    let prefix = parse_macro_input!(attr);
238    Visitor { prefix }.visit_item_mut(&mut item);
239    quote! {#item}.into()
240}
hax_rust_engine_macros/lib.rs

hax_rust_engine_macros/
lib.rs
